Merge branch 'dev0.8.0' of github.com:fastnlp/fastNLP into dev0.8.0

fastNLP/core/callbacks/callback.py

@@ -264,7 +264,7 @@ class Callback:
         r"""
         ``callback`` 的名称，我们会使用该名称从 ``checkpoint`` 中读取的相应的 ``state`` 并传递给 :meth:`on_load_checkpoint` 函数。
 
-        :return: 返回用于区分该 ``callback`` 实例的名称；
+        :return: 用于区分该 ``callback`` 实例的名称；
         """
         return self.__class__.__name__
 

fastNLP/core/callbacks/topk_saver.py

@@ -51,7 +51,6 @@ class Saver:
 
         self.timestamp_path = self.folder.joinpath(os.environ[FASTNLP_LAUNCH_TIME])
 
-    @rank_zero_call
     def save(self, trainer, folder_name):
         """
         执行保存的函数，将数据保存在::
@@ -62,10 +61,11 @@ class Saver:
 
         :param trainer: Trainer 对象
         :param folder_name: 保存的 folder 名称，将被创建。
-        :return: 返回实际发生保存的 folder 绝对路径。如果为 None 则没有创建。
+        :return: 实际发生保存的 folder 绝对路径。如果为 None 则没有创建。
         """
         folder = self.timestamp_path.joinpath(folder_name)
         folder.mkdir(parents=True, exist_ok=True)
+
         save_fn = getattr(trainer, self.save_fn_name)
         save_fn(
             folder=folder,
@@ -217,7 +217,7 @@ class TopkSaver(ResultsMonitor, Saver):
         self.topk_queue = TopkQueue(topk)
         self.save_evaluate_results = save_evaluate_results
 
-    @rank_zero_call
+    # 注意这里我们为了支持 torch_fsdp 去除了 ''@rank_zero_call''；
     def save_topk(self, trainer, results: Dict) -> Optional[str]:
         """
         根据 ``results`` 是否满足 topk 的相关设定决定是否保存，如果发生了保存，将返回保存的文件夹。如果返回为 ``None`` ，则说明此次没有满足

fastNLP/core/callbacks/utils.py

@@ -14,7 +14,7 @@ def _get_monitor_value(monitor: Union[callable, str], real_monitor: Optional[str
     :param monitor:
     :param real_monitor:
     :param res:
-    :return: 返回两个值（str, value)，其中str就是最终要到的key，value就是这个key对应的value。如果value为None说明当前results中没有
+    :return: 两个值（str, value)，其中str就是最终要到的key，value就是这个key对应的value。如果value为None说明当前results中没有
         找到对应的 monitor
     """
     if len(res) == 0 or monitor is None:

fastNLP/core/collators/collator.py

@@ -211,7 +211,7 @@ class Collator:
             若 ``pad_val`` 为 ``None`` ，该值无意义 。
         :param pad_fn: 指定当前 field 的 pad 函数，传入该函数则 ``pad_val``, ``dtype``, ``backend`` 等参数失效。``pad_fn`` 的输入为当前 field 的
             batch 形式。 Collator 将自动 unbatch 数据，然后将各个 field 组成各自的 batch 。
-        :return: 返回 Collator 自身；
+        :return: Collator 自身；
         """
         self._renew()
 
@@ -298,7 +298,7 @@ class Collator:
         :param field_names: field_name: 需要调整的 field 的名称。如果 :meth:`Dataset.__getitem__` 方法返回的是字典类型，则可以直接使用对应的
             field 的 key 来表示，如果是嵌套字典，可以使用元组表示多层次的 key，例如 ``{'a': {'b': 1}}`` 中可以使用 ``('a', 'b')``;
             如果 :meth:`Dataset.__getitem__` 返回的是 Sequence 类型，则可以使用 ``'_0'``, ``'_1'`` 表示序列中第 **0** 或 **1** 个元素。
-        :return: 返回 Collator 自身；
+        :return: Collator 自身；
         """
         self._renew()
         input_field_names = [(field, field) if isinstance(field, tuple) else ((field,), field)

fastNLP/core/controllers/evaluator.py

@@ -257,7 +257,7 @@ class Evaluator:
                ``metric_indicator_name#metric_name#dataloader_name``，其中 metric_indicator_name 可能不存在；
 
         :param num_eval_batch_per_dl: 每个 dataloader 测试前多少个 batch 的数据，-1 为测试所有数据。
-        :return: 返回评测得到的结果，是一个没有嵌套的字典；
+        :return: 评测得到的结果，是一个没有嵌套的字典；
         """
         assert isinstance(num_eval_batch_per_dl, int), "num_eval_batch_per_dl must be of int type."
         assert num_eval_batch_per_dl > 0 or num_eval_batch_per_dl == -1, "num_eval_batch_per_dl must be -1 or larger than 0."

fastNLP/core/controllers/trainer.py

@@ -57,13 +57,14 @@ class Trainer(TrainerEventTrigger):
 
     :param driver: 训练模型所使用的具体的驱动模式，应当为以下选择中的一个：``["auto", "torch", "paddle", "jittor", "fairscale", "deepspeed", "oneflow"]``：
         
-        1. 值为 ``"auto"`` 时，**FastNLP** 会根据传入模型的类型自行判断使用哪一种模式；
+        1. 值为 ``"auto"`` 时，**fastNLP** 会根据传入模型的类型自行判断使用哪一种模式；
         2. 其值为 ``"torch"`` 时，表示使用 :class:`~fastNLP.core.drivers.TorchSingleDriver` 或者 :class:`~fastNLP.core.drivers.TorchDDPDriver`；
-        3. 其值为 ``"paddle"`` 时，表示使用 :class:`~fastNLP.core.drivers.PaddleSingleDriver` 或者 :class:`~fastNLP.core.drivers.PaddleFleetDriver`；
-        4. 其值为 ``"jittor"`` 时，表示使用 :class:`~fastNLP.core.drivers.JittorSingleDriver` 或者 :class:`~fastNLP.core.drivers.JittorMPIDriver`；
-        5. 其值为 ``"fairscale"`` 时，表示使用 :class:`~fastNLP.core.drivers.FairScaleDriver`；
-        6. 其值为 ``"deepspeed"`` 时，表示使用 :class:`~fastNLP.core.drivers.DeepSpeedDriver`；
-        7. 其值为 ``"oneflow"`` 时，表示使用 :class:`~fastNLP.core.drivers.OneflowSingleDriver` 或者 :class:`~fastNLP.core.drivers.OneflowDDPDriver`；
+        3. 其值为 ``"torch_fsdp"`` 时，表示使用 :class:`~fastNLP.core.drivers.TorchFSDPDriver`；
+        4. 其值为 ``"paddle"`` 时，表示使用 :class:`~fastNLP.core.drivers.PaddleSingleDriver` 或者 :class:`~fastNLP.core.drivers.PaddleFleetDriver`；
+        5. 其值为 ``"jittor"`` 时，表示使用 :class:`~fastNLP.core.drivers.JittorSingleDriver` 或者 :class:`~fastNLP.core.drivers.JittorMPIDriver`；
+        6. 其值为 ``"fairscale"`` 时，表示使用 :class:`~fastNLP.core.drivers.FairScaleDriver`；
+        7. 其值为 ``"deepspeed"`` 时，表示使用 :class:`~fastNLP.core.drivers.DeepSpeedDriver`；
+        8. 其值为 ``"oneflow"`` 时，表示使用 :class:`~fastNLP.core.drivers.OneflowSingleDriver` 或者 :class:`~fastNLP.core.drivers.OneflowDDPDriver`；
         
         在指定了框架的情况下，具体使用哪一种取决于参数 ``device`` 的设置；
 
@@ -73,6 +74,10 @@ class Trainer(TrainerEventTrigger):
             这意味着当您传入一个 ``Driver`` 实例时，您传入给 ``Trainer`` 的 ``model`` 参数将会被忽略；也就是说模型在训练时使用的真正的模型是
             您传入的 ``Driver`` 实例中的模型；
 
+        .. note::
+
+            如果您选择使用 :mod:`deepspeed` 、:mod:`fairscale` 或 :mod:`torch.distributed.fsdp` 进行训练，请不要将 ``driver`` 的值设为 ``'auto'`` 。
+
     :param train_dataloader: 训练数据集，注意其必须是单独的一个数据集，不能是 :class:`List` 或者 :class:`Dict`；
 
         .. warning::
@@ -297,11 +302,26 @@ class Trainer(TrainerEventTrigger):
     :kwargs:
         * *torch_kwargs* -- ``TorchDriver`` 所需的其它参数，详见 :class:`~fastNLP.core.drivers.torch_driver.TorchSingleDriver` 和
           :class:`~fastNLP.core.drivers.torch_driver.TorchDDPDriver`；
+
+            .. note::
+
+                注意如果对于 ``TorchDDPDriver`` 中初始化 ``DistributedDataParallel`` 时有特别的参数，您可以通过在 ``torch_kwargs`` 中传入
+                ``ddp_kwargs`` 来实现，例如：
+
+                .. code-block::
+
+                    trainer = Trainer(
+                        ...,
+                        torch_kwargs = {'ddp_kwargs': {'find_unused_parameters': True, ...}}
+                    )
+
+                对于 ``TorchFSDPDriver`` 也是类似，只是对应的 ``**_kwargs`` 修改为 ``fsdp_kwargs``；
+
         * *paddle_kwargs* -- ``PaddleDriver`` 所需的其它参数，详见 :class:`~fastNLP.core.drivers.paddle_driver.PaddleSingleDriver` 和
           :class:`~fastNLP.core.drivers.paddle_driver.PaddleSingleDriver`；
         * *fairscale_kwargs* -- ``FairScaleDriver`` 所需的其它参数，详见 :class:`~fastNLP.core.drivers.torch_driver.FairScaleDriver`；
         * *deepspeed_kwargs* -- ``DeepSpeedDriver`` 所需的其它参数，详见 :class:`~fastNLP.core.drivers.torch_driver.DeepSpeedDriver`；
-        * *torch_kwargs* -- ``OneflowDriver`` 所需的其它参数，详见 :class:`~fastNLP.core.drivers.oneflow_driver.OneflowSingleDriver` 和
+        * *oneflow_kwargs* -- ``OneflowDriver`` 所需的其它参数，详见 :class:`~fastNLP.core.drivers.oneflow_driver.OneflowSingleDriver` 和
           :class:`~fastNLP.core.drivers.oneflow_driver.OneflowDDPDriver`；
         * *data_device* -- 一个具体的 driver 实例中，有 ``model_device`` 和 ``data_device``，前者表示模型所在的设备，后者表示
          当 ``model_device`` 为 None 时应当将数据迁移到哪个设备；
@@ -885,7 +905,7 @@ class Trainer(TrainerEventTrigger):
         :param marker: 用来标记该 callback 函数属于哪几个具体的 trainer 实例；两个特殊情况：1.当 ``marker`` 为 None（默认情况）时，
          表示该 callback 函数只属于代码下方最近的一个 trainer 实例；2.当 ``marker`` 为 'all' 时，该 callback 函数会被所有的 trainer
          实例使用；
-        :return: 返回原函数；
+        :return: 原函数；
         """
 
         def wrapper(fn: Callable) -> Callable:
@@ -1001,7 +1021,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def driver(self):
         """
-        :return: 返回 ``trainer`` 中的 ``driver`` 实例；
+        :return:  ``trainer`` 中的 ``driver`` 实例；
         """
         return self._driver
 
@@ -1012,7 +1032,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def train_batch_loop(self):
         """
-        :return: 返回 ``trainer`` 中的 ``train_batch_loop`` 实例；
+        :return:  ``trainer`` 中的 ``train_batch_loop`` 实例；
         """
         return self._train_batch_loop
 
@@ -1275,7 +1295,7 @@ class Trainer(TrainerEventTrigger):
             ``trainer.backward / zero_grad / step`` 函数的作用类似；
 
         :param batch: 一个 batch 的数据；
-        :return: 返回模型的前向传播函数所返回的结果；
+        :return: 模型的前向传播函数所返回的结果；
         """
         with self.driver.auto_cast():
             outputs = self.driver.model_call(batch, self._train_step, self._train_step_signature_fn)
@@ -1328,7 +1348,7 @@ class Trainer(TrainerEventTrigger):
         用来从用户模型的输出对象中抽取 ``loss`` 对象；
         目前支持 `outputs` 对象为 ``dict`` 或者 ``dataclass``；
 
-        :return: 返回被抽取出来的 ``loss`` 对象，例如如果是 ``pytorch``，那么返回的就是一个 tensor；
+        :return: 被抽取出来的 ``loss`` 对象，例如如果是 ``pytorch``，那么返回的就是一个 tensor；
         """
         if isinstance(outputs, Dict):
             try:
@@ -1375,7 +1395,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def n_epochs(self) -> int:
         r"""
-        :return: 返回当前训练的总体的 epoch 的数量；
+        :return: 当前训练的总体的 epoch 的数量；
         """
         return self.trainer_state.n_epochs
 
@@ -1386,7 +1406,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def cur_epoch_idx(self) -> int:
         r"""
-        :return: 返回当前正在第几个 epoch；
+        :return: 当前正在第几个 epoch；
         """
         return self.trainer_state.cur_epoch_idx
 
@@ -1397,7 +1417,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def global_forward_batches(self) -> int:
         """
-        :return: 返回从训练开始到当前总共训练了多少 batch 的数据；
+        :return: 从训练开始到当前总共训练了多少 batch 的数据；
         """
         return self.trainer_state.global_forward_batches
 
@@ -1408,7 +1428,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def batch_idx_in_epoch(self) -> int:
         r"""
-        :return: 返回在从当前的这个 epoch 开始，到现在共训练了多少 batch 的数据；
+        :return: 在从当前的这个 epoch 开始，到现在共训练了多少 batch 的数据；
         """
         return self.trainer_state.batch_idx_in_epoch
 
@@ -1419,7 +1439,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def num_batches_per_epoch(self) -> int:
         r"""
-        :return: 返回每一个 epoch 实际会训练多少个 batch 的数据；
+        :return: 每一个 epoch 实际会训练多少个 batch 的数据；
         """
         return self.trainer_state.num_batches_per_epoch
 
@@ -1430,7 +1450,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def n_batches(self) -> int:
         r"""
-        :return: 返回整体的训练中实际会训练多少个 batch 的数据；
+        :return: 整体的训练中实际会训练多少个 batch 的数据；
         """
         return self.trainer_state.n_batches
 
@@ -1443,7 +1463,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def model_device(self):
         r"""
-        :return: 返回当前模型所在的设备；注意该值在当且仅当在少数情况下为 ``None``，例如当使用 ``pytorch`` 时，仅当用户自己初始化 ``init_progress_group`` 时
+        :return: 当前模型所在的设备；注意该值在当且仅当在少数情况下为 ``None``，例如当使用 ``pytorch`` 时，仅当用户自己初始化 ``init_progress_group`` 时
             ``model_device`` 才为 None；
         """
         return self.driver.model_device
@@ -1451,7 +1471,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def data_device(self):
         r"""
-        :return: 返回数据会被迁移到的目的设备；
+        :return: 数据会被迁移到的目的设备；
         """
         return self.driver.data_device
 
@@ -1460,7 +1480,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def train_dataloader(self):
         """
-        :return: 返回用户传入的 ``train_dataloader``，注意该 ``dataloader`` 与用户传入给 ``Trainer`` 的 ``dataloader`` 对象是同一个对象，而我们在
+        :return: 用户传入的 ``train_dataloader``，注意该 ``dataloader`` 与用户传入给 ``Trainer`` 的 ``dataloader`` 对象是同一个对象，而我们在
         实际训练过程中使用的 ``dataloader`` 的状态可能有所更改；
         """
         return self._train_dataloader
@@ -1472,7 +1492,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def evaluate_dataloaders(self):
         """
-        :return: 返回用户传入的 ``evaluate_dataloaders``；
+        :return: 用户传入的 ``evaluate_dataloaders``；
         """
         return self._evaluate_dataloaders
 

fastNLP/core/controllers/utils/state.py

@@ -65,7 +65,7 @@ class TrainerState:
 
     def state_dict(self) -> Dict:
         r"""
-        :return: 返回用于断点重训来保存的状态字典；
+        :return: 用于断点重训来保存的状态字典；
         """
         return {"cur_epoch_idx": self.cur_epoch_idx, "global_forward_batches": self.global_forward_batches,
                 "batch_idx_in_epoch": self.batch_idx_in_epoch}

fastNLP/core/dataloaders/jittor_dataloader/fdl.py

@@ -146,7 +146,7 @@ class JittorDataLoader:
             若 ``pad_val`` 为 ``None`` ，该值无意义 。
         :param pad_fn: 指定当前 field 的 pad 函数，传入该函数则 ``pad_val``, ``dtype``, ``backend`` 等参数失效。``pad_fn`` 的输入为当前 field 的
             batch 形式。 collator 将自动 unbatch 数据，然后将各个 field 组成各自的 batch 。
-        :return: 返回使用的 collator
+        :return: 使用的 collator
         """
         collator = self._get_collator()
         if isinstance(collator, Collator):
@@ -177,7 +177,7 @@ class JittorDataLoader:
         :param field_names: field_name: 需要调整的 field 的名称。如果 :meth:`Dataset.__getitem__` 方法返回的是字典类型，则可以直接使用对应的
             field 的 key 来表示，如果是嵌套字典，可以使用元组表示多层次的 key，例如 ``{'a': {'b': 1}}`` 中可以使用 ``('a', 'b')``;
             如果 :meth:`Dataset.__getitem__` 返回的是 Sequence 类型，则可以使用 ``'_0'``, ``'_1'`` 表示序列中第 **0** 或 **1** 个元素。
-        :return: 返回使用的 collator
+        :return: 使用的 collator
         """
         collator = self._get_collator()
         if isinstance(collator, Collator):

fastNLP/core/dataloaders/oneflow_dataloader/fdl.py

@@ -167,7 +167,7 @@ class OneflowDataLoader(DataLoader):
             若 ``pad_val`` 为 ``None`` ，该值无意义 。
         :param pad_fn: 指定当前 field 的 pad 函数，传入该函数则 ``pad_val``, ``dtype``, ``backend`` 等参数失效。``pad_fn`` 的输入为当前 field 的
             batch 形式。 collator 将自动 unbatch 数据，然后将各个 field 组成各自的 batch 。
-        :return: 返回使用的 collator
+        :return: 使用的 collator
         """
         collator = self._get_collator()
         if isinstance(collator, Collator):
@@ -198,7 +198,7 @@ class OneflowDataLoader(DataLoader):
         :param field_names: field_name: 需要调整的 field 的名称。如果 :meth:`Dataset.__getitem__` 方法返回的是字典类型，则可以直接使用对应的
             field 的 key 来表示，如果是嵌套字典，可以使用元组表示多层次的 key，例如 ``{'a': {'b': 1}}`` 中可以使用 ``('a', 'b')``;
             如果 :meth:`Dataset.__getitem__` 返回的是 Sequence 类型，则可以使用 ``'_0'``, ``'_1'`` 表示序列中第 **0** 或 **1** 个元素。
-        :return: 返回使用的 collator
+        :return: 使用的 collator
         """
         collator = self._get_collator()
         if isinstance(collator, Collator):

fastNLP/core/dataloaders/paddle_dataloader/fdl.py

@@ -199,7 +199,7 @@ class PaddleDataLoader(DataLoader):
             若 ``pad_val`` 为 ``None`` ，该值无意义 。
         :param pad_fn: 指定当前 field 的 pad 函数，传入该函数则 ``pad_val``, ``dtype``, ``backend`` 等参数失效。``pad_fn`` 的输入为当前 field 的
             batch 形式。 collator 将自动 unbatch 数据，然后将各个 field 组成各自的 batch 。
-        :return: 返回使用的 collator
+        :return: 使用的 collator
         """
         collator = self._get_collator()
         if isinstance(collator, Collator):
@@ -230,7 +230,7 @@ class PaddleDataLoader(DataLoader):
         :param field_names: field_name: 需要调整的 field 的名称。如果 :meth:`Dataset.__getitem__` 方法返回的是字典类型，则可以直接使用对应的
             field 的 key 来表示，如果是嵌套字典，可以使用元组表示多层次的 key，例如 ``{'a': {'b': 1}}`` 中可以使用 ``('a', 'b')``;
             如果 :meth:`Dataset.__getitem__` 返回的是 Sequence 类型，则可以使用 ``'_0'``, ``'_1'`` 表示序列中第 **0** 或 **1** 个元素。
-        :return: 返回使用的 collator
+        :return: 使用的 collator
         """
         collator = self._get_collator()
         if isinstance(collator, Collator):

fastNLP/core/dataloaders/torch_dataloader/fdl.py

@@ -167,7 +167,7 @@ class TorchDataLoader(DataLoader):
             若 ``pad_val`` 为 ``None`` ，该值无意义 。
         :param pad_fn: 指定当前 field 的 pad 函数，传入该函数则 ``pad_val``, ``dtype``, ``backend`` 等参数失效。``pad_fn`` 的输入为当前 field 的
             batch 形式。 collator 将自动 unbatch 数据，然后将各个 field 组成各自的 batch 。
-        :return: 返回使用的 collator
+        :return: 使用的 collator
         """
         collator = self._get_collator()
         if isinstance(collator, Collator):
@@ -198,7 +198,7 @@ class TorchDataLoader(DataLoader):
         :param field_names: field_name: 需要调整的 field 的名称。如果 :meth:`Dataset.__getitem__` 方法返回的是字典类型，则可以直接使用对应的
             field 的 key 来表示，如果是嵌套字典，可以使用元组表示多层次的 key，例如 ``{'a': {'b': 1}}`` 中可以使用 ``('a', 'b')``;
             如果 :meth:`Dataset.__getitem__` 返回的是 Sequence 类型，则可以使用 ``'_0'``, ``'_1'`` 表示序列中第 **0** 或 **1** 个元素。
-        :return: 返回使用的 collator
+        :return: 使用的 collator
         """
         collator = self._get_collator()
         if isinstance(collator, Collator):

fastNLP/core/dataloaders/torch_dataloader/mix_dataloader.py

@@ -99,35 +99,24 @@ class _MixCollateFn:
 
 class MixDataLoader(DataLoader):
     """
-    针对以下四种情况提供的 ``MixDataLoader``， 目前只支持 ``torch`` 框架的版本， 其中 mode 的取值范围为 ``['sequential', 'mix', 'polling', "Sampler"]``:
+    针对以下四种情况提供的 ``MixDataLoader``， 目前只支持 **pytorch** 框架的版本， 其中 mode 的取值范围为 ``['sequential', 'mix', 'polling', 'Sampler']``:
 
         * 当 mode 为 ``'sequential'`` 时，``MixDataLoader``  将 ``datasets`` 的序列或者字典视为一个混合大数据集， 按照 datasets 数据集序列或者字典的顺序一个
           接一个的 sample 完所有数据。
         * 当 mode 为 ``'mix'`` 时， ``MixDataLoader``  将 ``datasets`` 的序列或者字典视为一个混合大数据集， 然后根据用户输入的 idx 序列随机 sample
           混合数据集 datasets 的数据组成一个 batch 序列返回。
         * 当 mode 为 ``'polling'`` 时， ``MixDataLoader`` 按照 ``datasets`` 数据集的顺序， 先从第一个数据集采样一个 batch 的数据返回，
-          再从第二数据集采样一个 batch 数据返回， 直至最后一个数据集采样一个 batch 数据返回后再从第一个数据采样第二个 batch 数据返回，直至所有的数据集都被轮询的采样完。
-        * 当 mode 为 ``"Sampler"`` 时， 该 Sampler 是实现 __iter__() 的实例化对象， 其功能是每次 iter 时返回一个 batch 序列， 其类型为 List[int];
+          再从第二数据集采样一个 batch 数据返回， 直至最后一个数据集采样一个 batch 数据返回后再从第一个数据采样第二个 batch 数据返回，直至所有的数据集都被轮询地的采样完。
+        * 当 mode 为 ``'Sampler'`` 时， 该 Sampler 是实现 __iter__() 的实例化对象， 其功能是每次 iter 时返回一个 batch 序列， 其类型为 List[int];
           且 Sampler 必须将输入的 datasets 视为一个混合大数据集， 其 index 范围为 ``0<idx<len(datasets[0])+...+len(datasets[x])``, 然后参数
           ``sampler``, ``drop_last``, ``ds_ratio`` 均无效。
 
     :param datasets: 实现了 __getitem__() 和 __len__() 对象的序列或者字典。
-    :param mode: mode 控制 ``MixDataLoader`` 运行模式。 mode 的取值范围为 ``['sequential', 'mix', 'polling', "Sampler"]``:
-
-        * 当 mode 为 ``'sequential'`` 时，``MixDataLoader``  将 datasets 的序列或者字典视为一个混合大数据集， 按照 datasets 数据集序列或者字典的顺序一个
-          接一个的 sample 完所有数据。
-        * 当 mode 为 ``'mix'`` 时， ``MixDataLoader``  将 datasets 的序列或者字典视为一个混合大数据集， 然后根据用户输入的 idx 序列随机sample
-          混合数据集 datasets 的数据组成一个 batch 序列返回。
-        * 当 mode 为 ``'polling'`` 时， ``MixDataLoader`` 按照 datasets 数据集的顺序， 先从第一个数据集采样一个 batch 的数据返回，
-          再从第二数据集采样一个 batch 数据返回， 直至最后一个数据集采样一个 batch 数据返回后再从第一个数据采样第二个 batch 数据返回，直至所有的数据集都被轮询的采样完。
-        * 当 mode 为 ``"Sampler"`` 时， 该 Sampler 是实现 __iter__() 的实例化对象， 其功能是每次 iter 时返回一个 batch 序列， 其类型为 List[int];
-          且 Sampler 必须将输入的 datasets 视为一个混合大数据集， 其 index 范围为 ``0<idx<len(datasets[0])+...+len(datasets[x])``, 然后参数
-          sampler, drop_last, ds_ratio 均无效。
-
+    :param mode: ``mode`` 控制 ``MixDataLoader`` 运行模式。 ``mode`` 的取值范围为 ``['sequential', 'mix', 'polling', 'Sampler']``，每种模式的详细功能见上文。
     :param collate_fn: 用于从 dataset 取到的一个 batch 数据进行打包处理的 Callable 函数。 其取值可以为 ``['auto', Callable, List[Callable], Dict[str, Callable]]``:
 
         * collate_fn 为 ``'auto'`` 时, ``MixDataLoader``  datasets 序列或者dict 初始化一个 :class:`~fastNLP.core.collators.Collator`  作为其默认值，
-          需要注意的是只有当 datasets 包含的所以 dataset 的数据都为 ``List`` 或者 ``Dict`` 类型时才能使用。否则只能用户自己定义 collate_fn .
+          需要注意的是只有当 datasets 包含的所以 dataset 的数据都为 ``List`` 或者 ``Dict`` 类型时才能使用。否则只能用户自己定义 collate_fn 。
         * collate_fn 为  :class:`Callable` 时， 该 collate_fn 会被 datasets 序列或者dict 的所有数据所共享。该 Callable 函数应当接受一个 batch 参数作为输入，
           batch 是一个 List 对象且 List 中的每一条数据都是 dataset 的一条数据；该 Callable 函数还应当返回一个对象。
         * collate_fn 为 ``Dict[str, Callable]`` 时， datasets 的 key 必须和 callable_fn 的 key 一致。 ``MixDataLoader`` 会将 ``collate_fn[key]``
@@ -135,12 +124,12 @@ class MixDataLoader(DataLoader):
 
 
     :param sampler: 实现了 __len__() 和 __iter__() 的实例化对象，其 __iter__() 方法每次都会返回 dataset 的一个下标 index ，其取值范围为
-        ``[None, str, Dict[str, "Sampler"]]``:
+        ``[None, str, Dict[str, Sampler]]``:
 
-        * sampler 为 ``None`` 时， ``MixDataLoader`` 默认初始化 ``torch`` 的 ``SequentialSampler`` 作为默认值。其功能时顺序返回 dataset 的下标。
-        * sampler 为 ``str`` 时， sampler 选择范围为 ``[rand, seq]``。当 sampler 为 ``rand`` 时，``MixDataLoader`` 默认初始化 ``torch`` 的  ``RandomSampler``
-          作为默认值， 其功能时随机采样 dataset 的下标并返回。 当 sampler 为 ``seq`` 时， ``MixDataLoader`` 默认初始化 ``torch`` 的 ``SequentialSampler`` 作为默认值。其功能时顺序返回 dataset 的下标。
-        * sampler 为 ``Dict[str, "Sampler"]`` 时， ``Sampler`` 为用户定义的实现了 __len__() 和 __iter__() 的实例化对象。 其每次 iter 必须返回一个 int 下标。
+        * sampler 为 ``None`` 时， ``MixDataLoader`` 默认初始化 **pytorch** 的 ``SequentialSampler`` 作为默认值。其功能时顺序返回 dataset 的下标。
+        * sampler 为 ``str`` 时， sampler 选择范围为 ``[rand, seq]``。当 sampler 为 ``rand`` 时，``MixDataLoader`` 默认初始化 **pytorch** 的  ``RandomSampler``
+          作为默认值， 其功能时随机采样 dataset 的下标并返回。 当 sampler 为 ``seq`` 时， ``MixDataLoader`` 默认初始化 **pytorch** 的 ``SequentialSampler`` 作为默认值。其功能时顺序返回 dataset 的下标。
+        * sampler 为 ``Dict[str, Sampler]`` 时， ``Sampler`` 为用户定义的实现了 __len__() 和 __iter__() 的实例化对象。 其每次 iter 必须返回一个 int 下标。
           Dict 的 str 必须和 datasets 的 key 一致。 也即是 ``Dict[str, Sampler]`` 为 datasets 字典的每个 dataset 初始化了一个 Sampler。
 
     :param num_workers: 当 ``num_workers > 0`` 时, ``MixDataLoader`` 会开启 ``num_workers`` 个子进程来处理数据， 可以加快数据处理速度，但同时

fastNLP/core/dataset/dataset.py

@@ -637,7 +637,7 @@ class DataSet:
 
         :param progress_desc: 如果不为 ``None``，则会显示当前正在处理的进度条的名称；
         :param progress_bar: 显示进度条的方式，支持 ``["rich", "tqdm", None]``。
-        :return: 返回一个字典
+        :return: 一个字典
         """
         assert len(self) != 0, "Null DataSet cannot use apply_field()."
         if not self.has_field(field_name=field_name):
@@ -762,7 +762,7 @@ class DataSet:
 
         :param progress_desc: 当 progress_bar 不为 ``None`` 时，可以显示当前正在处理的进度条名称
         :param progress_bar: 显示进度条的方式，支持 ``["rich", "tqdm", None]``。
-        :return: 返回一个字典
+        :return: 一个字典
         """
         assert callable(func), "The func is not callable."
         assert len(self) != 0, "Null DataSet cannot use apply()."
@@ -1013,7 +1013,7 @@ class DataSet:
             若 ``pad_val`` 为 ``None`` ，该值无意义 。
         :param pad_fn: 指定当前 field 的 pad 函数，传入该函数则 ``pad_val``, ``dtype``, ``backend`` 等参数失效。``pad_fn`` 的输入为当前 field 的
             batch 形式。 Collator 将自动 unbatch 数据，然后将各个 field 组成各自的 batch 。
-        :return: 返回自身的 collator；
+        :return: 自身的 collator；
         """
         if isinstance(self.collator, Collator):
             self.collator.set_pad(field_name=field_name, pad_val=pad_val, dtype=dtype, pad_fn=pad_fn, backend=backend)
@@ -1031,7 +1031,7 @@ class DataSet:
         :param field_names: field_name: 需要调整的 field 的名称。如果 :meth:`Dataset.__getitem__` 方法返回的是字典类型，则可以直接使用对应的
             field 的 key 来表示，如果是嵌套字典，可以使用元组表示多层次的 key，例如 ``{'a': {'b': 1}}`` 中可以使用 ``('a', 'b')``;
             如果 :meth:`Dataset.__getitem__` 返回的是 Sequence 类型，则可以使用 ``'_0'``, ``'_1'`` 表示序列中第 **0** 或 **1** 个元素。
-        :return: 返回自身的 collator；
+        :return: 自身的 collator；
         """
         if isinstance(self.collator, Collator):
             self.collator.set_ignore(*field_names)

fastNLP/core/drivers/__init__.py

@@ -3,6 +3,8 @@ __all__ = [
     'TorchDriver',
     "TorchSingleDriver",
     "TorchDDPDriver",
+    "FairScaleDriver",
+    "TorchFSDPDriver",
     "DeepSpeedDriver",
     "PaddleDriver",
     "PaddleSingleDriver",
@@ -10,8 +12,6 @@ __all__ = [
     "JittorDriver",
     "JittorSingleDriver",
     "JittorMPIDriver",
-    'TorchSingleDriver',
-    'TorchDDPDriver',
     'PaddleDriver',
     'PaddleSingleDriver',
     'PaddleFleetDriver',
@@ -27,7 +27,8 @@ __all__ = [
     'optimizer_state_to_device'
 ]
 
-from .torch_driver import TorchDriver, TorchSingleDriver, TorchDDPDriver, DeepSpeedDriver, torch_seed_everything, optimizer_state_to_device
+from .torch_driver import TorchDriver, TorchSingleDriver, TorchDDPDriver, DeepSpeedDriver, FairScaleDriver, \
+    TorchFSDPDriver, torch_seed_everything, optimizer_state_to_device
 from .jittor_driver import JittorDriver, JittorMPIDriver, JittorSingleDriver
 from .paddle_driver import PaddleDriver, PaddleFleetDriver, PaddleSingleDriver, paddle_seed_everything
 from .oneflow_driver import OneflowDriver, OneflowSingleDriver, OneflowDDPDriver, oneflow_seed_everything

fastNLP/core/drivers/choose_driver.py

@@ -3,15 +3,17 @@ from typing import Union, Optional, List
 from .driver import Driver
 from ..utils import is_torch_module, is_paddle_module, is_jittor_module, is_oneflow_module
 
+__all__ = []
 
 def choose_driver(model, driver: Union[str, Driver], device: Optional[Union[int, List[int], str]], **kwargs) -> Driver:
     r"""
-    根据输入的参数 'gpus' 的格式来决定具体的工作模式;
+    根据输入的参数 ``driver`` 和 ``device`` 的格式来决定具体的工作模式。
 
-    :param model: 运行过程中使用的具体的最原始的模型；
-    :param driver: 应当为字符串或者 `Driver` 实例，表示运行中具体使用的训练/评测模式；
-    :param device: 具体的形式请参见 `fastNLP.core.drivers.torch_driver.utils.initialize_torch_dirver` 的注释；
-    :param kwargs: 其余的传给 `Driver` 的参数；
+    :param model: 运行过程中使用的具体的最原始的模型。
+    :param driver: 训练模型所使用的具体的驱动模式，应当为以下选择中的一个：``["auto", "torch", "paddle", "jittor", "fairscale", "deepspeed", "oneflow", "torch_fsdp"]``，分别对应
+        各种框架。值为 ``'auto'`` 时，将会根据模型的类型进行选择。
+    :param device: 训练使用的设备。详细的格式可以查阅 :class:`~fastNLP.core.controllers.Trainer` 中的说明。
+    :param kwargs: 其余的传给 `Driver` 的参数。
     """
 
     # 如果用户直接传进来一个 driver 实例，我们就直接返回回去，目前用户需要自己保证传进来的 driver 的正确性；
@@ -30,7 +32,7 @@ def choose_driver(model, driver: Union[str, Driver], device: Optional[Union[int,
         else:
             raise ValueError(f"Cannot choose driver automatically based on model, please set `driver` specifically.")
 
-    if driver in {"torch", "fairscale", "deepspeed"}:
+    if driver in {"torch", "fairscale", "deepspeed", "torch_fsdp"}:
         from fastNLP.core.drivers.torch_driver.initialize_torch_driver import initialize_torch_driver
         return initialize_torch_driver(driver, device, model, **kwargs)
     elif driver in {"jittor"}:

fastNLP/core/drivers/driver.py

@@ -18,9 +18,10 @@ from fastNLP.core.utils import nullcontext
 class Driver(ABC):
     r"""
     用来初始化 `Driver` 的基类，所有定制的 `driver` 都需要继承此类；
-    fastNLP 提供的 driver 实例都会同时被 Trainer 和 Evaluator 调用；
-    :param model: 训练或者评测的模型，需要注意该模型可能为用户已经使用类似 `torch.nn.DataParallel` 或者
-        `torch.nn.parallel.DistributedDataParallel` 包裹过的模型；
+    **fastNLP** 提供的 driver 实例都会同时被 :class:`~fastNLP.core.controllers.Trainer` 和 :class:`~fastNLP.core.controllers.Evaluator` 调用。
+
+    :param model: 训练或者评测的模型，需要注意该模型可能为用户已经使用类似 :class:`torch.nn.DataParallel` 或者
+        :class:`torch.nn.parallel.DistributedDataParallel` 包裹过的模型。
     """
 
     def __init__(self, model):
@@ -33,29 +34,32 @@ class Driver(ABC):
     @abstractmethod
     def setup(self):
         r"""
-        该函数用来初始化训练环境，例如将模型迁移到对应的设备上等；
+        该函数用来初始化训练环境，例如将模型迁移到对应的设备上等。
         多卡的 ``driver`` 的该函数要更为复杂一些，例如其可能需要开启多进程之间的通信环境，以及设置一些环境变量和其余所需要的变量值；
         """
 
     def set_dist_repro_dataloader(self, dataloader, dist=None, reproducible: bool = False):
         r"""
-        根据输入的 ``dataloader`` 得到一个 支持分布式 （``distributed``） 与 可复现的 (``reproducible``) 的 dataloader。
-
-        :param dataloader: 根据 ``dataloader`` 设置其对应的分布式版本以及可复现版本；
-        :param dist: 应当为一个字符串，其值应当为以下之一：``[None, "dist", "unrepeatdist"]``；为 ``None`` 时，表示不需要考虑当前 dataloader
-            切换为分布式状态；为 ``dist`` 时，表示该 dataloader 应该保证每个 gpu 上返回的 batch 的数量是一样多的，允许出现少量 sample ，在
-            不同 gpu 上出现重复；为 ``unrepeatdist`` 时，表示该 dataloader 应该保证所有 gpu 上迭代出来的数据合并起来应该刚好等于原始的
-            数据，允许不同 gpu 上 batch 的数量不一致。
-            其中 trainer 中 kwargs 的参数 ``use_dist_sampler`` 为 ``True`` 时，该值为 ``dist``；
-            否则为 ``None``，evaluator 中的 kwargs 的参数 ``use_dist_sampler`` 为 ``True`` 时，该值为 ``unrepeatdist``，否则为 ``None``；
-            注意当 dist 为 ReproducibleSampler, ReproducibleBatchSampler 时，是断点重训加载时 driver.load_checkpoint 函数在调用；
-            当 dist 为 str 或者 None 时，是 trainer 在初始化时调用该函数；
+        根据输入的 ``dataloader`` 得到一个 支持分布式 （**distributed**） 与 可复现的 (**reproducible**) 的 dataloader。
+
+        :param dataloader: 根据 ``dataloader`` 设置其对应的分布式版本以及可复现版本。
+        :param dist: 应当为一个字符串，其值应当为以下之一：``[None, "dist", "unrepeatdist"]``，并且根据在 :class:`~fastNLP.core.controllers.Trainer`
+            和 :class:`~fastNLP.core.controllers.Evaluator` 中 *kwargs* 的参数 ``use_dist_sampler`` 和调用时机不同，对应不同的值：
+        
+            * 当 ``use_dist_sampler`` 为 ``False`` ，且在 :class:`~fastNLP.core.controllers.Trainer` 或 :class:`~fastNLP.core.controllers.Evaluator`
+              **初始化** 中被调用时，参数值为 ``None`` ，表示不需要考虑当前 ``dataloader`` 切换为分布式状态；
+            * 当 ``use_dist_sampler`` 为 ``True`` ，且在 :class:`~fastNLP.core.controllers.Trainer` **初始化** 中被调用时，参数值为 ``"dist"`` ，表示该
+              ``dataloader`` 应该保证每个 gpu 上返回的 batch 的数量是一样多的，允许出现少量 sample 在不同 gpu 上出现重复；
+            * 当 ``use_dist_sampler`` 为 ``True`` ，且在 :class:`~fastNLP.core.controllers.Evaluator` **初始化** 中被调用时，参数值为 ``"unrepeatdist"`` ，
+              表示该 ``dataloader`` 应该保证所有 gpu 上迭代出来的数据合并起来应该刚好等于原始的数据，允许不同 gpu 上 batch 的数量不一致；
+            * 当 **断点重训加载** 中调用 :meth:`load_checkpoint` 时，该函数也会被调用，且 ``dist`` 值为 :class:`~fastNLP.core.samplers.ReproducibleSampler`
+              或 :class:`~fastNLP.core.samplers.ReproducibleBatchSampler` ，此时表示需要用 ``dist`` 代表的 sampler 或 batch_sampler 重新实例化一个新的 dataloader；
 
         :param reproducible: 如果为 ``False``，不要做任何考虑；如果为 ``True``，需要保证返回的 dataloader 可以保存当前的迭代状态，使得
-            该状态可以加载到一个全新的 dataloader 中然后恢复其状态；
-        :return: 应当返回一个被替换 sampler 后的新的 dataloader 对象 (注意此处一定需要返回一个新的 dataloader 对象) ；此外，
-            如果传入的 dataloader 中是 ReproducibleSampler 或者 ReproducibleBatchSampler 需要重新初始化一个放入返回的
-            dataloader 中。如果 dist 为空，且 reproducible 为 False，可直接返回原对象。
+            该状态可以加载到一个全新的 dataloader 中然后恢复其状态。
+        :return: 应当返回一个被替换 sampler 后的 **新的** dataloader 对象 (注意此处一定需要返回一个新的 dataloader 对象) ；此外，
+            如果传入的 ``dataloader`` 中是 :class:`~fastNLP.core.samplers.ReproducibleSampler` 或者 :class:`~fastNLP.core.samplers.ReproducibleBatchSampler`
+            需要 **重新初始化** 一个放入返回的 dataloader 中。如果 ``dist`` 为空，且 ``reproducible`` 为 ``False``，可直接返回原对象。
         """
         if dist is None and reproducible is False:
             return dataloader
@@ -64,64 +68,69 @@ class Driver(ABC):
 
     def set_deterministic_dataloader(self, dataloader):
         r"""
-        为了确定性训练要对 ``dataloader`` 进行修改，保证在确定随机数种子后，每次重新训练得到的结果是一样的；例如对于 ``pytorch`` 的 ``dataloader``，其
-        需要将 ``worker_init_fn`` 替换；
+        为了确定性训练要对 ``dataloader`` 进行修改，保证在确定随机数种子后，每次重新训练得到的结果是一样的；例如对于 **pytorch** 的 ``dataloader``，其
+        需要将 ``worker_init_fn`` 替换。
         """
 
     def set_sampler_epoch(self, dataloader, cur_epoch_idx):
         r"""
-        对于分布式的 ``sampler``，例如 ``pytorch`` 的 ``DistributedSampler``，其需要在每一个 ``epoch`` 前设置随机数种子，来保证每一个进程上的 ``shuffle`` 是一样的；
+        对于分布式的 ``sampler``，例如 **pytorch** 的 :class:`DistributedSampler`，其需要在每一个 ``epoch`` 前设置随机数种子，来保证每一个进程上的 ``shuffle`` 是一样的；
         ``dataloader`` 中可能真正发挥作用的是 ``batch_sampler`` 也可能是 ``sampler``。
 
-        :param dataloader: 需要设置 ``epoch`` 的 ``dataloader``；
-        :param cur_epoch_idx: 当前是第几个 ``epoch``；
+        :param dataloader: 需要设置 ``epoch`` 的 ``dataloader``
+        :param cur_epoch_idx: 当前是第几个 ``epoch``
         """
 
     @abstractmethod
     def model_call(self, batch, fn: Callable, signature_fn: Optional[Callable]) -> Dict:
         r"""
         通过调用 ``fn`` 来实现训练时的前向传播过程；
-        注意 ``Trainer`` 和 ``Evaluator`` 会调用该函数来实现网络的前向传播过程，其中传入该函数的参数 ``fn`` 是函数 ``get_model_call_fn`` 所返回的
-        函数；
+        注意 :class:`~fastNLP.core.controllers.Trainer` 和 :class:`~fastNLP.core.controllers.Evaluator` 会调用该函数来
+        实现网络的前向传播过程，其中传入该函数的参数 ``fn`` 是函数 :meth:`get_model_call_fn` 所返回的函数。
 
-        :param batch: 当前的一个 batch 的数据；可以为字典或者其它类型；
+        :param batch: 当前的一个 batch 的数据；可以为字典或者其它类型。
         :param fn: 调用该函数进行一次计算。
-        :param signature_fn: 由 ``Trainer`` 传入的用于网络前向传播一次的签名函数，因为当 batch 是一个 ``Dict`` 的时候，我们会自动调用 ``auto_param_call`` 函
-            数，而一些被包裹的模型需要暴露其真正的函数签名，例如 ``DistributedDataParallel`` 的调用函数是 ``forward``，但是需要其函数签名为 ``model.module.forward``；
-        :return: 返回由 ``fn`` 返回的结果（应当为一个 ``dict`` 或者 ``dataclass``，但是不需要我们去检查）；
+        :param signature_fn: 由 :class:`~fastNLP.core.controllers.Trainer` 传入的用于网络前向传播一次的签名函数，因为当
+            batch 是一个 :class:`Dict` 的时候，我们会自动调用 :func:`fastNLP.core.utils.auto_param_call` 函数，而一些被
+            包裹的模型需要暴露其真正的函数签名，例如 :class:`DistributedDataParallel` 的调用函数是 ``forward``，但是需要其
+            函数签名为 ``model.module.forward``。
+        :return: 由 ``fn`` 返回的结果（应当为一个 :class:`dict` 或者 :class:`dataclass` ，但是不需要我们去检查）。
         """
         raise NotImplementedError("Each specific driver should implemented its own `model_call` function.")
 
     @abstractmethod
     def get_model_call_fn(self, fn: str) -> Tuple:
         r"""
-        该函数会接受 ``Trainer`` 的 ``train_fn`` 或者 ``Evaluator`` 的 ``evaluate_fn``，返回一个实际用于调用 ``driver.model_call`` 时传入的函数参数；
-        该函数会在 ``Trainer`` 和 ``Evaluator`` 在 ``driver.setup`` 函数之后调用；
+        该函数会接受 :class:`~fastNLP.core.controllers.Trainer` 的 ``train_fn`` 或者 :class:`~fastNLP.core.controllers.Evaluator`
+        的 ``evaluate_fn``，返回一个实际用于调用 :meth:`model_call` 时传入的函数参数；该函数会由 :class:`~fastNLP.core.controllers.Trainer`
+        和 :class:`~fastNLP.core.controllers.Evaluator` 在 :func:`driver.setup` 函数之后调用。
 
-        之所以设置该函数的目的在于希望将具体的 model_call function 从 driver 中抽离出来，然后将其附着在 Trainer 或者 Evaluator 身上；
+        之所以设置该函数的目的在于希望将具体的 model_call function 从 driver 中抽离出来，然后将其附着在 ``Trainer`` 或者 ``Evaluator`` 身上；
         这样是因为在新版的设计中，使用 model 的哪种方法来进行 ``train step`` 或者 ``evaluate step`` 是通过额外的参数 ``train_fn`` 和
-        ``evaluate_fn`` 来确定的，而二者又分别是通过 Trainer 和 Evaluator 来控制的；因此不能将确定具体的 ``train step fn`` 和
-        ``evaluate step fn`` 的逻辑放在每一个 driver 的初始化的时候（因此在 Trainer 初始化第一个 driver 时，Evaluator 还没有初始化，但是
-        ``evaluate step fn`` 的确定却需要 Evaluator 的初始化），因此我们将这一逻辑抽象到这一函数当中；
+        ``evaluate_fn`` 来确定的，而二者又分别是通过 ``Trainer`` 和 ``Evaluator`` 来控制的；因此不能将确定具体的 ``train step fn`` 和
+        ``evaluate step fn`` 的逻辑放在每一个 driver 的初始化的时候（因此在 ``Trainer`` 初始化第一个 driver 时，``Evaluator`` 还没有初始化，但是
+        ``evaluate step fn`` 的确定却需要 Evaluator 的初始化），因此我们将这一逻辑抽象到这一函数当中.
 
         这一函数应当通过参数 ``fn`` 来判断应当返回的实际的调用的函数，具体逻辑如下所示：
-        1. 如果 fn == "train_step" or "evaluate_step"，那么对传入的模型进行检测，如果模型没有定义方法 ``fn``，则默认调用模型的 ``forward``
-        函数，然后给出 warning；
-        2. 如果 fn 是其他字符串，那么如果模型没有定义方法 ``fn`` 则直接报错；
 
-        注意不同的 driver 需要做额外的检测处理，例如在 DDPDriver 中，当传入的模型本身就是 DistributedDataParallel 中，我们只能调用模型的
-        forward 函数，因此需要额外的 warning；这一点特别需要注意的问题在于 driver 自己在 setup 时也会对模型进行改变（DDPDriver），因此
-        可能需要额外标记最初传入 driver 的模型是哪种形式的；
+            1. 如果 ``fn`` == "train_step" or "evaluate_step"，那么对传入的模型进行检测，如果模型没有定义方法 ``fn``，则默认调用模型的 :meth:`forward`
+               函数，然后给出 warning；
+            2. 如果 ``fn`` 是其他字符串，那么如果模型没有定义方法 ``fn`` 则直接报错；
+
+        注意不同的 driver 需要做额外的检测处理，例如在 :class:`~fastNLP.core.drivers.torch_driver.TorchDDPDriver` 中，当传入的模型本身就是
+        :class:`DistributedDataParallel` 时，我们只能调用模型的 :meth:`forward` 函数，因此需要额外的 warning；这一点特别需要注意的问题在于
+        driver 自己在 setup 时也会对模型进行改变（ :class:`~fastNLP.core.drivers.torch_driver.TorchDDPDriver` ），因此可能需要额外标记最初
+        传入 driver 的模型是哪种形式的.
 
-        :param fn: 应当为一个字符串，该函数通过该字符串判断要返回模型的哪种方法；
-        :return: 返回一个元组，包含两个函数，用于在调用 driver.model_call 时传入；
+        :param fn: 一个字符串，该函数通过该字符串判断要返回模型的哪种方法
+        :return: 一个元组，包含两个函数，用于在调用 :meth:`model_call` 时传入
         """
         raise NotImplementedError("Each specific driver should implemented its own `get_model_call_fn` function.")
 
     @property
     def model(self):
         r"""
-        :return: 返回 driver 中在实际训练或者评测时所使用的模型；
+        :return:  driver 中在实际训练或者评测时所使用的模型。
         """
         return self._model
 
@@ -132,8 +141,8 @@ class Driver(ABC):
     @property
     def optimizers(self) -> List:
         r"""
-        如下所示，driver 返回的 optimizers 一定是一个 List，如果用户直接向 Trainer 传入一个单独的 optimizer，我们会使用一个 List 将其
-        包裹；
+        如下所示，driver 返回的 :attr:`optimizers` 一定是一个 :class:`List`，如果用户直接向 :class:`~fastNLP.core.controllers.Trainer` 传入一个单独的 optimizer，
+        我们会使用一个 List 将其包裹；
 
         :return: List[optimizer0, optimizer1, optimizer2, ...]
         """
@@ -150,7 +159,7 @@ class Driver(ABC):
     @property
     def model_device(self):
         r"""
-        :return: 返回 driver 中模型实际所在的设备；
+        :return: driver 中模型实际所在的设备。
         """
         return self._model_device
 
@@ -161,55 +170,54 @@ class Driver(ABC):
     @property
     def data_device(self):
         """
-        :return: 返回 driver 中数据默认会被迁移到的设备；
+        :return: driver 中数据默认会被迁移到的设备。
         """
         return self.model_device
 
     @staticmethod
     def _check_optimizer_legality(optimizers):
         r"""
-        对于用户传入 trainer 的每一个 optimizer，检测其是否合理，因为不同的深度学习框架所使用的的 optimizer 是不相同的；
+        对于用户传入 trainer 的每一个 optimizer，检测其是否合理，因为不同的深度学习框架所使用的的 optimizer 是不相同的。
 
-        :param optimizers: 需要检测的 `optimizers`；
+        :param optimizers: 需要检测的 `optimizers`。
         """
         raise NotImplementedError(
             "Each specific driver should implemented its own `_check_optimizer_legality` function.")
 
     def check_dataloader_legality(self, dataloader):
         """
-        检测 DataLoader 是否合法，如果不合法，会 raise TypeError 。
+        检测 ``dataloader`` 是否合法，如果不合法，会 ``raise TypeError`` 。
 
         :param dataloder:
-        :return:
         """
 
     def set_optimizers(self, optimizers=None):
         r"""
-        trainer 会调用该函数将用户传入的 optimizers 挂载到 driver 实例上；
+        trainer 会调用该函数将用户传入的 ``optimizers`` 挂载到 driver 实例上。
         """
         self.optimizers = optimizers
 
     @abstractmethod
     def backward(self, loss):
         r"""
-        实现深度学习中的反向传播过程；
+        实现深度学习中的反向传播过程。
 
-        :param loss: 用来实现反向传播的损失函数值；
+        :param loss: 用来实现反向传播的损失函数值
         """
         raise NotImplementedError("Each specific driver should implemented its own `backward` function.")
 
     @abstractmethod
     def step(self):
         r"""
-        实现深度学习中的参数的优化更新过程，应当直接通过优化器 optimizers 来更新参数；
+        实现深度学习中的参数的优化更新过程，应当直接通过优化器 :attr:`optimizers` 来更新参数。
         """
         raise NotImplementedError("Each specific driver should implemented its own `step` function.")
 
     @abstractmethod
     def zero_grad(self):
         r"""
-        实现深度学习中的梯度的置零操作，应当直接通过优化器 optimizers 来将梯度置零；
-        注意梯度累积不需要在这里实现，trainer 已经在内部实现了梯度累积；
+        实现深度学习中的梯度的置零操作，应当直接通过优化器 :attr:`optimizers` 来将梯度置零；
+        注意梯度累积不需要在这里实现，trainer 已经在内部实现了梯度累积。
 
         """
         raise NotImplementedError("Each specific driver should implemented its own `zero_grad` function.")
@@ -217,26 +225,26 @@ class Driver(ABC):
     def get_model_no_sync_context(self):
         r"""
         返回一个用于关闭多进程之间 model 中的自动互相同步操作的 context 上下文对象；只有多卡的 driver 需要单独实现该函数，
-        单卡的 driver 不需要；
+        单卡的 driver 不需要。
 
-        :return: 返回一个类似于 DistributedDataParallel(model).no_sync 的 context 上下文对象；
+        :return: 一个类似于 ``DistributedDataParallel(model).no_sync`` 的 context 上下文对象
         """
         return nullcontext
 
     def get_evaluate_context(self):
         r"""
-        返回一个不计算梯度的环境用来对模型进行评测；
+        返回一个不计算梯度的环境用来对模型进行评测。
 
-        :return: 一个类似 `torch.no_grad` 的 context 上下文对象；
+        :return: 一个类似 ``torch.no_grad`` 的 context 上下文对象
         """
         return nullcontext
 
     @property
     def auto_cast(self):
         r"""
-        fp16 的上下文环境；
+        fp16 的上下文环境。
 
-        :return: 返回一个用于 fp16 计算的上下文环境；
+        :return: 一个用于 fp16 计算的上下文环境
         """
         return self._auto_cast
 
@@ -247,20 +255,19 @@ class Driver(ABC):
     @abstractmethod
     def save_model(self, filepath: Union[str, Path, BytesIO], only_state_dict: bool = True, **kwargs):
         r"""
-        保存模型的函数；注意函数 `save` 是用来进行断点重训的函数；
+        保存模型的函数；注意函数 :meth:`save_checkpoint` 是用来进行断点重训的函数。
 
-        :param filepath: 保存文件的文件位置（需要包括文件名）或一个 BytesIO 对象；
-        :param only_state_dict: 是否只保存模型的 `state_dict`；
-        :param model_save_fn: 用户传入的用来代替该函数本身保存逻辑的函数；如果该参数不为 None，那么我们会调用 model_save_fn(path)；
+        :param filepath: 保存文件的文件位置（需要包括文件名）或一个 BytesIO 对象
+        :param only_state_dict: 是否只保存模型的 `state_dict`
         """
         raise NotImplementedError("Each specific driver should implemented its own `save_model` function.")
 
     @abstractmethod
     def load_model(self, filepath: Union[str, Path, BytesIO], only_state_dict: bool = False, **kwargs):
         r"""
-        加载模型的函数；将 filepath 中的模型加载并赋值给当前 model 。
+        加载模型的函数；将 ``filepath`` 中的模型加载并赋值给当前 model 。
 
-        :param filepath: 需要被加载的对象的文件位置（需要包括文件名）或一个 ``BytesIO`` 对象；
+        :param filepath: 需要被加载的对象的文件位置（需要包括文件名）或一个 ``BytesIO`` 对象。
         :param load_state_dict: 保存的文件是否只是模型的权重，还是完整的模型。即便是保存的完整的模型，此处也只能使用尝试加载filepath
             模型中的权重到自身模型，而不会直接替代当前 Driver 中的模型。
         """
@@ -269,18 +276,18 @@ class Driver(ABC):
     @abstractmethod
     def save_checkpoint(self, folder, states: Dict, dataloader, only_state_dict: bool = True, should_save_model: bool = True,
              **kwargs):
+
         r"""
-        断点重训的保存函数，该函数会负责保存模型和 optimizers, fp16 的 state_dict；以及模型的保存（若 should_save_model 为 True）
+        断点重训的保存函数，该函数会负责保存优化器、fp16 状态和 sampler 的状态，以及模型的保存（若 ``should_save_model`` 为 ``True``）
 
-        :param folder: 保存断点重训的状态的文件夹；save_checkpoint 函数应该在下面新增两（一）个文件 的 FASTNLP_CHECKPOINT_FILENAME 文件与
-            FASTNLP_MODEL_FILENAME （如果 should_save_model 为 True ）。把 model 相关的内容放入到 FASTNLP_MODEL_FILENAME 文件
-            中，将传入的 states 以及自身产生其它状态一并保存在 FASTNLP_CHECKPOINT_FILENAME 里面。
-        :param states: 由 trainer 传入的一个字典，其中已经包含了为了实现断点重训所需要保存的其它对象的状态，Driver 应该只需要保存
-            该对象即可， Driver 应该不需要理解该对象，同时在 driver.load_checkpoint() 的时候，需要将 states 返回回去，load_checkpoint() 返回的值与这里的
-            传入的值保持一致。
+        :param folder: 保存断点重训的状态的文件夹；:meth:`save_checkpoint` 函数应该在该路径下面下面新增名为 ``FASTNLP_CHECKPOINT_FILENAME`` 与
+            ``FASTNLP_MODEL_FILENAME`` （如果 ``should_save_model`` 为 ``True`` ）的文件。把 model 相关的内容放入到 ``FASTNLP_MODEL_FILENAME`` 文件
+            中，将传入的 ``states`` 以及自身产生的其它状态一并保存在 ``FASTNLP_CHECKPOINT_FILENAME`` 里面。
+        :param states: 由 :class:`~fastNLP.core.controllers.Trainer` 传入的一个字典，其中已经包含了为了实现断点重训所需要保存的其它对象的状态。Driver 应该
+            只需要保存该对象而不需要理解该对象，同时在 :meth:`load_checkpoint` 的时候需要将 ``states`` 返回回去，返回的值与这里传入的值保持一致。
         :param dataloader: 正在使用的 dataloader，需要保存里面的状态使得之后可以从当前迭代的位置恢复。
-        :param only_state_dict: 是否只保存模型的参数，当 should_save_model 为 False ，该参数无效。
-        :param should_save_model: 是否应该保存模型，如果为False，Driver 将不负责 model 的保存。
+        :param only_state_dict: 是否只保存模型的参数，当 ``should_save_model`` 为 ``False`` ，该参数无效。
+        :param should_save_model: 是否应该保存模型，如果为 ``False`` ，Driver 将不负责 model 的保存。
         """
         raise NotImplementedError("Each specific driver should implemented its own `save_checkpoint` function.")
 
@@ -288,113 +295,115 @@ class Driver(ABC):
     def load_checkpoint(self, folder: Union[str, Path], dataloader, only_state_dict: bool = True, should_load_model: bool = True,
              **kwargs) -> Dict:
         r"""
-        断点重训的加载函数，注意该函数会负责读取数据，并且恢复 optimizers , fp16 的 state_dict 和 模型（根据 should_load_model ）和；
-        其它在 Driver.save_checkpoint() 函数中执行的保存操作，然后将一个 state 字典返回给 trainer （ 内容为Driver.save_checkpoint() 接受到的 states ）。
-
+        断点重训的加载函数，该函数会负责读取数据，并且恢复优化器 、sampler 的状态和模型（如果 ``should_load_model`` 为 True）以及其它在 :meth:`save_checkpoint`
+        函数中执行的保存操作，然后将一个 state 字典返回给 :class:`~fastNLP.core.controllers.Trainer` （ 内容为 :meth:`save_checkpoint` 接受到的 ``states`` ）。
+        
         该函数应该在所有 rank 上执行。
 
-        :param folder: 读取该 folder 下的 FASTNLP_CHECKPOINT_FILENAME 文件与 FASTNLP_MODEL_FILENAME
+        :param folder: 读取该 folder 下的 ``FASTNLP_CHECKPOINT_FILENAME`` 文件与 ``FASTNLP_MODEL_FILENAME``
             （如果 should_load_model 为True）。
-        :param dataloader: 当前给定 dataloader，需要根据保存的 dataloader 状态合理设置。若该值为 None ，是不需要返回 'dataloader'
-            以及 'batch_idx_in_epoch' 这两个值。
-        :param only_state_dict: 读取的，当 should_save_model 为 False ，该参数无效。如果为 True ，说明保存的内容为权重；如果为
+        :param dataloader: 当前给定 dataloader，需要根据保存的 dataloader 状态合理设置。若该值为 ``None`` ，则不需要返回 ``'dataloader'``
+            以及 ``'batch_idx_in_epoch'`` 这两个值。
+        :param only_state_dict: 是否仅读取模型的 state_dict ，当 ``should_save_model`` 为 ``False`` ，该参数无效。如果为 ``True`` ，说明保存的内容为权重；如果为
             False 说明保存的是模型，但也是通过当前 Driver 的模型去加载保存的模型的权重，而不是使用保存的模型替换当前模型。
-        :param should_load_model: 是否应该加载模型，如果为False，Driver 将不负责加载模型。若该参数为 True ，但在保存的状态中没有
+        :param should_load_model: 是否应该加载模型，如果为 ``False`` ，Driver 将不负责加载模型。若该参数为 ``True`` ，但在保存的状态中没有
             找到对应的模型状态，则报错。
-        :return: 需要返回 save_checkpoint 函数输入的 states 内容
+        :return: :meth:`save_checkpoint` 函数输入的 ``states`` 内容。除此之外，还返回的内容有：
+
+            * *dataloader* -- 根据传入的 ``dataloader`` 与读取出的状态设置为合理状态的 dataloader。在当前 ``dataloader`` 样本数与读取出的 sampler 样本数
+              不一致时报错。
+            * *batch_idx_in_epoch* -- :class:`int` 类型的数据，表明当前 epoch 进行到了第几个 batch 。请注意，该值不能仅通过保存的数据中读取的，因为前后两次运行的
+              ``batch_size`` 可能有变化，而应该符合以下等式::
 
-            * *dataloader* -- 返回的是根据传入的 dataloader 与 保存的状态一起设置为合理的状态，可以返回的对象与传入的dataloader是同一个。
-             在保存与当前传入 data sample 数目不一致时报错。
+                返回的 dataloader 还会产生的 batch 数量 + batch_idx_in_epoch = 原来不断点训练时的 batch 的总数
+              
+              由于 ``返回的 dataloader 还会产生的batch数`` 在 ``batch_size`` 与 ``drop_last`` 参数给定的情况下，无法改变，因此只能通过调整 ``batch_idx_in_epoch``
+              这个值来使等式成立。一个简单的计算原则如下：
 
-            * *batch_idx_in_epoch* -- int 类型的数据，表明当前 epoch 进行到了进行到了第几个 batch 了。 请注意，该值不能是只能通过保存的
-             数据中读取的，因为前后两次运行 batch_size 可能由变化。该数字的原则应该符合以下等式
-             '返回 dataloader 还会产生的batch数量' + 'batch_idx_in_epoch' = '原来不断点训练的batch的总数' 。
-             由于 '返回 dataloader 还会产生的batch数量' 这个数量在 batch_size 与 drop_last 参数给定的情况下，无法改变，因此
-             只能通过调整 batch_idx_in_epoch 这个值来使等式成立。一个简单的计算原则如下
-             当drop_last为True，等同于 floor(sample_in_this_rank/batch_size) - floor(num_left_samples/batch_size);
-             当drop_last为False，等同于 ceil(sample_in_this_rank/batch_size) - ceil(num_left_samples/batch_size)。
+                * drop_last 为 ``True`` 时，等同于 floor(sample_in_this_rank/batch_size) - floor(num_left_samples/batch_size)；
+                * drop_last 为 ``False`` 时，等同于 ceil(sample_in_this_rank/batch_size) - ceil(num_left_samples/batch_size)。
         """
-        raise NotImplementedError("Each specific driver should implemented its own `load_checkpoint` function.")
 
     @staticmethod
     def tensor_to_numeric(tensor, reduce: Optional[str] = None):
         r"""
         将一个 ``tensor`` 对象（仅处理当前 driver 使用的 tensor 即可）转换为 python 的 ``numeric`` 对象；如果 ``tensor`` 只包含一个
-        元素则返回 ``float`` 或 ``int``；
+        元素则返回 ``float`` 或 ``int``。
 
-        :param tensor: 需要被转换的 `tensor` 对象；
+        :param tensor: 需要被转换的 ``tensor`` 对象
         :param reduce: 可选 ``['sum', 'max', 'mea', 'min']``，如果不为 ``None`` 将使用该 ``reduce`` 方法来处理当前 ``tensor`` 再返回
-            ``float`` 或 ``int`` 对象；
-        :return: 转换后返回的结果；
+            :class:`float` 或 :class:`int` 对象
+        :return: 转换后返回的结果
         """
         raise NotImplementedError("Each specific driver should implemented its own `tensor_to_numeric` function.")
 
     @abstractmethod
     def set_model_mode(self, mode: str):
         r"""
-        设置模型为 `train` / `eval` 的模式；目的是为切换模型训练和推理（会关闭dropout等）模式；
+        设置模型为 ``train`` 或 ``eval`` 的模式；目的是为切换模型的训练和推理（会关闭 dropout 等）模式。
 
-        :param mode: 应为二者之一：["train", "eval"]；
+        :param mode: 应为二者之一：``["train", "eval"]``
         """
 
     def unwrap_model(self):
         r"""
         保证用户拿到的模型一定是最原始的模型；
-        注意因为我们把保存模型的主要逻辑和代码移到了 `Driver` 中，因此在 `save_model` 函数中，一定要先调用此函数来保证我们保存的模型一定是
+        注意因为我们把保存模型的主要逻辑和代码移到了 `Driver` 中，因此在 :meth:`save_model` 函数中，一定要先调用此函数来保证我们保存的模型一定是
         最为原始的模型；
-        需要注意用户本身传入的模型就是经过类似 `torch.nn.DataParallel` 或者 `torch.nn.parallel.DistributedDataParallel` 包裹的模型，
-        因此在该函数内需要先判断模型的类别；
+        需要注意用户本身传入的模型就是经过类似 :class:`torch.nn.DataParallel` 或者 :class:`torch.nn.parallel.DistributedDataParallel` 包裹的模型，
+        因此在该函数内需要先判断模型的类别。
 
-        :return: 返回最原始的模型，例如没有被 `DistributedDataParallel` 包裹的模型；
+        :return: 最原始的模型，例如没有被 :class:`DistributedDataParallel` 包裹的模型。
         """
 
     @staticmethod
     def move_model_to_device(model, device):
         r"""
-        用来将模型转移到指定的 device 上；
-        之所以写成 `staticmethod`，是因为一方面在 `Driver` 中我们要使用 `unwrap_model` 来拿到最原始的模型，另一方面，在 `save_model`
-        中，我们需要先将模型移到 cpu 后，又再移到 gpu 上，因此不适宜在该函数内部调用 `unwrap_model`，而是将 model 作为该函数的参数；
+        用来将模型转移到指定的 ``device`` 上；
+        之所以写成 :class:`staticmethod`，是因为一方面在 `Driver` 中我们要使用 :meth:`unwrap_model` 来拿到最原始的模型，另一方面，在 :meth`save_model`
+        中，我们需要先将模型移到 cpu 后，又再移到 gpu 上，因此不适宜在该函数内部调用 :meth:`unwrap_model`，而是将 ``model`` 作为该函数的参数。
         """
 
     @abstractmethod
     def move_data_to_device(self, batch):
         r"""
-        将数据迁移到指定的机器上；batch 可能是 list 也可能 dict ，或其嵌套结构；
+        将数据迁移到指定的机器上；``batch`` 是包含了张量的数据集合，可以是 **List**、**Dict** 等嵌套类型。
 
-        :return: 将移动到指定机器上的 batch 对象返回；
+        :return: 移动到指定机器上的 ``batch`` 对象
         """
 
     def get_local_rank(self) -> int:
         r"""
-        返回当前的local_rank，本函数的返回值只在运行分布式训练的时候有实际含义；
+        返回当前的 ``local_rank``，本函数的返回值只在运行分布式训练的时候有实际含义。
 
-        :return: 一个整数值，表示当前进程在当前这台机器上的序号；
+        :return: 一个整数值，表示当前进程在当前这台机器上的序号
         """
         return 0
 
     def barrier(self):
         r"""
         用于在多进程工作时同步各进程的工作进度，运行快的进程运行到这里会等待运行慢的进程，只有所有进程都运行到此函数时，所有的进程才会继续运行；
-        仅在多分布式训练场景中有使用；
+        仅在多分布式训练场景中有使用。
 
-        注意，该函数的行为会受到 FASTNLP_NO_SYNC 的影响。仅当 FASTNLP_NO_SYNC 在 os.environ 中不存在，或小于 1 时才真的执行 barrier；
+        注意，该函数的行为会受到环境变量 ``FASTNLP_NO_SYNC`` 的影响。仅当 ``FASTNLP_NO_SYNC`` 在 ``os.environ`` 中不存在，或小于 **1** 时
+        才真的执行 :meth:`barrier`。
         """
 
     def is_distributed(self) -> bool:
         r"""
-        当前的 driver 实例是否是分布式的；
+        当前的 driver 实例是否是分布式的。
 
-        :return: 返回一个 bool 值，如果当前的 driver 实例是用于分布式的，那么返回 True；
+        :return: 一个 bool 值，如果当前的 driver 实例是用于分布式的，那么返回 ``True``
         """
         return False
 
     def on_exception(self):
         r"""
-        该函数用于在训练或者预测过程中出现错误时正确地关掉其它的进程，这一点是通过在多进程 driver 调用 open_subprocess 的时候将每一个进程
-        的 pid 记录下来，然后在出现错误后，由出现错误的进程手动地将其它进程 kill 掉；
+        该函数用于在训练或者预测过程中出现错误时正确地关掉其它的进程，这一点是通过在多进程 driver 调用 :meth:`open_subprocess` 的时候将每一个进程
+        的 pid 记录下来，然后在出现错误后，由出现错误的进程手动地将其它进程 kill 掉。
 
-        因此，每一个多进程 driver 如果想要该函数能够正确地执行，其需要在自己的 open_subprocess（开启多进程的函数）中正确地记录每一个进程的
-        pid 的信息；
+        因此，每一个多进程 driver 如果想要该函数能够正确地执行，其需要在自己的 :meth:`open_subprocess` （开启多进程的函数）中正确地记录每一个进程的
+        pid 的信息；单卡 driver 不需要这个函数。
         """
         # 单卡 driver 不需要这个函数；
         if self._pids is not None:
@@ -419,10 +428,10 @@ class Driver(ABC):
         从 ``src`` 端将 ``obj`` 对象（可能是 ``tensor``，可能是 ``object`` ）broadcast 到其它所有进程。如果是非 ``tensor`` 的对象会尝试使用 ``pickle`` 进行打包进行
         传输，然后再 ``dst`` 处再加载回来。仅在分布式的 ``driver`` 中有实际意义。
 
-        :param obj: obj，可能是 ``Tensor`` 或 嵌套类型的数据；
-        :param src: source 的 ``global rank``；
-        :param group: 所属的通信组；
-        :return: 输入的 ``obj``；
+        :param obj: obj，可能是 ``Tensor`` 或 嵌套类型的数据
+        :param src: source 的 ``global rank``
+        :param group: 所属的通信组
+        :return: 输入的 ``obj``
         """
         if not self.is_distributed():
             return obj
@@ -431,12 +440,12 @@ class Driver(ABC):
 
     def all_gather(self, obj, group) -> List:
         r"""
-        将 obj 互相传送到其它所有的 rank 上，其中 obj 可能是 Tensor，也可能是嵌套结构的 object 。如果不是基础类型的数据，尝试通过
+        将 ``obj`` 互相传送到其它所有的 rank 上，其中 ``obj`` 可能是 Tensor，也可能是嵌套结构的 object 。如果不是基础类型的数据，尝试通过
         pickle 进行序列化，接收到之后再反序列化。
 
-        :param obj: 可以是 ``float/int/bool/np.ndarray/{}/[]/Tensor`` 等；
-        :param group: 用于不同进程之间互相通信的通信组；
-        :return: 返回值应该是 ``[obj0, obj1, ...]``，其中 ``obj1`` 是 ``rank0`` 上的对象，``obj1`` 是 ``rank1`` 上的对象；
+        :param obj: 可以是 ``float/int/bool/np.ndarray/{}/[]/Tensor`` 等类型的数据
+        :param group: 用于不同进程之间互相通信的通信组
+        :return: 返回值应该是 ``[obj0, obj1, ...]``，其中 ``obj0`` 是 ``rank0`` 上的对象，``obj1`` 是 ``rank1`` 上的对象。以此类推
         """
         if not self.is_distributed():
             return [obj]

fastNLP/core/drivers/jittor_driver/initialize_jittor_driver.py

@@ -18,9 +18,9 @@ def initialize_jittor_driver(driver: str, device: Union[str, int, List[int]], mo
 
         创建多卡的 driver
 
-    :param driver: 该参数的值应为以下之一：``["jittor"]``；
-    :param device: ``jittor`` 运行的设备；
-    :param model: 训练或者评测的具体的模型；
+    :param driver: 该参数的值应为以下之一：``["jittor"]``
+    :param device: ``jittor`` 运行的设备
+    :param model: 训练或者评测的具体的模型
     :param kwargs: 
 
     :return: :class:`~fastNLP.core.JittorSingleDriver` 或 :class:`~fastNLP.core.JittorMPIDriver` 实例；

fastNLP/core/drivers/jittor_driver/jittor_driver.py

@@ -40,19 +40,22 @@ __all__ = [
 
 class JittorDriver(Driver):
     r"""
-    ``Jittor`` 框架的 ``Driver``，是 ``JittorSingleDevice`` 和 ``JittorMPIDriver`` 的父类。
+    实现了 **jittor** 框架训练功能的基本 ``Driver``。这个类被以下子类继承：
+    
+        1. :class:`~fastNLP.core.drivers.jittor_driver.JittorSingleDriver` ：实现了使用单卡和 ``cpu`` 训练的具体功能；
+        2. :class:`~fastNLP.core.drivers.jittor_driver.JittorMPIDriver` ：实现了使用 ``mpi`` 启动 **jittor** 分布式训练的功能；
 
     .. warning::
 
         您不应当直接初始化该类，然后传入给 ``Trainer``，换句话说，您应当使用该类的子类 ``JittorSingleDriver`` 和 ``TorchDDPDriver``，而不是
-        该类本身；
+        该类本身。
 
     .. note::
 
-        您可以在使用 ``JittorSingleDevice`` 和 ``JittorMPIDriver`` 时使用 ``JittorDriver`` 提供的接口；
+        您可以在使用 ``JittorSingleDriver`` 和 ``JittorMPIDriver`` 时使用 ``JittorDriver`` 提供的接口。
 
-    :param model: 训练时使用的 **jittor** 模型；
-    :param fp16: 是否开启混合精度训练;
+    :param model: 训练时使用的 **jittor** 模型
+    :param fp16: 是否开启混合精度训练
     :param jittor_kwargs:
     """
     def __init__(self, model, fp16: bool = False, jittor_kwargs: Dict = None, **kwargs):
@@ -73,6 +76,11 @@ class JittorDriver(Driver):
         self.wo_auto_param_call = kwargs.get("model_wo_auto_param_call", False)
 
     def check_dataloader_legality(self, dataloader):
+        """
+        检测 DataLoader 是否合法。支持的类型包括 :class:`~fastNLP.core.dataloaders.JittorDataLoader`、 :class:`jittor.dataset.Dataset` 。
+
+        :param dataloder:
+        """
         if not isinstance(dataloader, (Dataset, JittorDataLoader, OverfitDataLoader)):
             raise TypeError(f"{Dataset} or {JittorDataLoader} is expected, instead of `{type(dataloader)}`")
         if len(dataloader) == 0:
@@ -87,14 +95,23 @@ class JittorDriver(Driver):
                                  f"not {type(each_optimizer)}.")
 
     def step(self):
+        r"""
+        实现参数的优化更新过程
+        """
         for optimizer in self.optimizers:
             optimizer.step()
 
     def backward(self, loss):
+        """
+        对 ``loss`` 进行反向传播
+        """        
         for optimizer in self.optimizers:
             optimizer.backward(loss)
 
     def zero_grad(self):
+        """
+        实现梯度置零的过程
+        """
         for optimizer in self.optimizers:
             optimizer.zero_grad()
 
@@ -102,8 +119,8 @@ class JittorDriver(Driver):
         r"""
         将模型保存到 ``filepath`` 中。
 
-        :param filepath: 保存文件的文件位置（需要包括文件名）；
-        :param only_state_dict: 在 **Jittor** 中，该参数无效，**Jittor** 仅支持保存模型的 ``state_dict``。
+        :param filepath: 保存文件的文件位置
+        :param only_state_dict: 在 **Jittor** 中，该参数无效，因为 **Jittor** 仅支持保存模型的 ``state_dict``。
         """
         if not only_state_dict:
             logger.rank_zero_warning(
@@ -119,7 +136,7 @@ class JittorDriver(Driver):
         r"""
         加载模型的函数；将 ``filepath`` 中的模型加载并赋值给当前 ``model`` 。
 
-        :param filepath: 保存文件的文件位置（需要包括文件名）；
+        :param filepath: 保存文件的文件位置
         :param load_state_dict: 在 **Jittor** 中，该参数无效，**Jittor** 仅支持加载模型的 ``state_dict``。
         """
         if not only_state_dict:
@@ -133,6 +150,17 @@ class JittorDriver(Driver):
         model.load(filepath)
 
     def save_checkpoint(self, folder: Path, states: Dict, dataloader, only_state_dict: bool = True, should_save_model: bool = True, **kwargs):
+        r"""
+        断点重训的保存函数，该函数会负责保存 **优化器** 和 **sampler** 的状态，以及 **模型** （若 ``should_save_model`` 为 ``True``）
+
+        :param folder: 保存断点重训的状态的文件夹；:meth:`save_checkpoint` 函数应该在该路径下面下面新增名为 ``FASTNLP_CHECKPOINT_FILENAME`` 与
+            ``FASTNLP_MODEL_FILENAME`` （如果 ``should_save_model`` 为 ``True`` ）的文件。把 model 相关的内容放入到 ``FASTNLP_MODEL_FILENAME`` 文件
+            中，将传入的 ``states`` 以及自身产生的其它状态一并保存在 ``FASTNLP_CHECKPOINT_FILENAME`` 里面。
+        :param states: 由 :class:`~fastNLP.core.controllers.Trainer` 传入的一个字典，其中已经包含了为了实现断点重训所需要保存的其它对象的状态。
+        :param dataloader: 正在使用的 dataloader。
+        :param only_state_dict: 是否只保存模型的参数，当 ``should_save_model`` 为 ``False`` ，该参数无效。
+        :param should_save_model: 是否应该保存模型，如果为 ``False`` ，Driver 将不负责 model 的保存。
+        """
         dataloader_args = self.get_dataloader_args(dataloader)
         if dataloader_args.sampler:
             sampler = dataloader_args.sampler
@@ -185,7 +213,36 @@ class JittorDriver(Driver):
         logger.debug("Load optimizer state dict.")
 
     def load_checkpoint(self, folder: Path, dataloader, only_state_dict: bool = True, should_load_model: bool = True, **kwargs) -> Dict:
-        
+        r"""
+        断点重训的加载函数，该函数会负责读取数据，并且恢复 **优化器** 、**sampler** 的状态和 **模型** （如果 ``should_load_model`` 为 True）以及其它
+        在 :meth:`save_checkpoint` 函数中执行的保存操作，然后将一个 state 字典返回给 :class:`~fastNLP.core.controllers.Trainer` （ 内容为 :meth:`save_checkpoint` 
+        接受到的 ``states`` ）。
+
+        该函数应该在所有 rank 上执行。
+
+        :param folder: 读取该 folder 下的 ``FASTNLP_CHECKPOINT_FILENAME`` 文件与 ``FASTNLP_MODEL_FILENAME``
+            （如果 should_load_model 为True）。
+        :param dataloader: 当前给定 dataloader，需要根据保存的 dataloader 状态合理设置。若该值为 ``None`` ，则不需要返回 ``'dataloader'``
+            以及 ``'batch_idx_in_epoch'`` 这两个值。
+        :param only_state_dict: 是否仅读取模型的 state_dict ，当 ``should_save_model`` 为 ``False`` ，该参数无效。如果为 ``True`` ，说明保存的内容为权重；如果为
+            False 说明保存的是模型，但也是通过当前 Driver 的模型去加载保存的模型的权重，而不是使用保存的模型替换当前模型。
+        :param should_load_model: 是否应该加载模型，如果为 ``False`` ，Driver 将不负责加载模型。若该参数为 ``True`` ，但在保存的状态中没有
+            找到对应的模型状态，则报错。
+        :return: :meth:`save_checkpoint` 函数输入的 ``states`` 内容。除此之外，还返回的内容有：
+
+            * *dataloader* -- 根据传入的 ``dataloader`` 与读取出的状态设置为合理状态的 dataloader。在当前 ``dataloader`` 样本数与读取出的 sampler 样本数
+              不一致时报错。
+            * *batch_idx_in_epoch* -- :class:`int` 类型的数据，表明当前 epoch 进行到了第几个 batch 。请注意，该值不能仅通过保存的数据中读取的，因为前后两次运行的
+              ``batch_size`` 可能有变化，而应该符合以下等式::
+
+                返回的 dataloader 还会产生的 batch 数量 + batch_idx_in_epoch = 原来不断点训练时的 batch 的总数
+              
+              由于 ``返回的 dataloader 还会产生的batch数`` 在 ``batch_size`` 与 ``drop_last`` 参数给定的情况下，无法改变，因此只能通过调整 ``batch_idx_in_epoch``
+              这个值来使等式成立。一个简单的计算原则如下：
+
+                * drop_last 为 ``True`` 时，等同于 floor(sample_in_this_rank/batch_size) - floor(num_left_samples/batch_size)；
+                * drop_last 为 ``False`` 时，等同于 ceil(sample_in_this_rank/batch_size) - ceil(num_left_samples/batch_size)。
+        """
         states = jt.load(str(folder.joinpath(FASTNLP_CHECKPOINT_FILENAME)))
 
         # 1. 加载 optimizers 的状态；
@@ -232,6 +289,11 @@ class JittorDriver(Driver):
         return states
 
     def get_evaluate_context(self):
+        r"""
+        返回一个不计算梯度的上下文环境用来对模型进行评测；
+
+        :return: 上下文对象 ``jittor.no_grad``
+        """
         return jt.no_grad
 
     @staticmethod
@@ -241,6 +303,12 @@ class JittorDriver(Driver):
         """
         ...
 
+    def move_data_to_device(self, batch: 'jt.Var'):
+        """
+        将数据迁移到指定的机器上；**jittor** 会自动为变量分配设备无需手动迁移，因此这个函数只是简单地返回 ``batch``。
+        """
+        return batch
+
     def move_data_to_device(self, batch):
         """
         将数据 ``batch`` 转移到指定的设备上。由于 **Jittor** 会自动为数据分配设备，因此该函数实际上无效。
@@ -250,11 +318,11 @@ class JittorDriver(Driver):
     @staticmethod
     def tensor_to_numeric(tensor, reduce=None):
         r"""
-        将一个 :class:`jittor.Var` 对象转换为 转换成 python 中的数值类型；
+        将一个 :class:`jittor.Var` 对象转换为 转换成 python 中的数值类型。
 
-        :param tensor: :class:`jittor.Var` 类型的对象；
-        :param reduce: 当 tensor 是一个多数值的张量时，应当使用何种归一化操作来转换成单一数值，应当为以下类型之一：``['max', 'min', 'sum', 'mean']``；
-        :return: 返回一个单一数值，其数值类型是 python 中的基本的数值类型，例如 ``int，float`` 等；
+        :param tensor: :class:`jittor.Var` 类型的对象
+        :param reduce: 当 tensor 是一个多数值的张量时，应当使用何种归一化操作来转换成单一数值，应当为以下类型之一：``['max', 'min', 'sum', 'mean']``。
+        :return: 返回一个单一数值，其数值类型是 python 中的基本的数值类型，例如 ``int，float`` 等
         """
         if tensor is None:
             return None
@@ -274,29 +342,45 @@ class JittorDriver(Driver):
         )
 
     def set_model_mode(self, mode: str):
+        r"""
+        设置模型为 ``train`` 或 ``eval`` 的模式；目的是为切换模型的训练和推理（会关闭 dropout 等）模式。
+
+        :param mode: 应为二者之一：``["train", "eval"]``
+        """
         assert mode in {"train", "eval"}
         getattr(self.model, mode)()
 
     @property
     def data_device(self):
-        return self.model_device
-
-    def move_data_to_device(self, batch: 'jt.Var'):
         """
-        **jittor** 暂时没有提供数据迁移的函数，因此这个函数只是简单地返回 **batch**
+        :return: 数据默认会被迁移到的设备
         """
-        return batch
+        return self.model_device
 
     def set_deterministic_dataloader(self, dataloader: Union["JittorDataLoader", "Dataset"]):
+        r"""
+        为了确定性训练要对 ``dataloader`` 进行修改，保证在确定随机数种子后，每次重新训练得到的结果是一样的。 **jittor** 暂时不提供
+        该功能。
+        """
         ...
 
     def set_sampler_epoch(self, dataloader: Union["JittorDataLoader", "Dataset"], cur_epoch_idx: int):
+        r"""
+        对于分布式的 ``sampler``，需要在每一个 ``epoch`` 前设置随机数种子，来保证每一个进程上的 ``shuffle`` 是一样的。
+
+        :param dataloader: 需要设置 ``epoch`` 的 ``dataloader``
+        :param cur_epoch_idx: 当前是第几个 ``epoch``
+        """
         # 保证 ddp 训练时的 shuffle=True 时的正确性，因为需要保证每一个进程上的 sampler 的shuffle 的随机数种子是一样的；
         if callable(getattr(dataloader.sampler, "set_epoch", None)):
             dataloader.sampler.set_epoch(cur_epoch_idx)
 
     @staticmethod
     def get_dataloader_args(dataloader: Union["JittorDataLoader", "Dataset"]):
+        """
+        从 ``dataloader`` 中获取参数 ``dataset``, ``batch_sampler``, ``sampler``, ``batch_size``, ``shuffle`` 
+        和 ``drop_last`` 。
+        """
         @dataclass
         class Res:
             dataset: Optional[Dataset] = None

fastNLP/core/drivers/jittor_driver/mpi.py

@@ -145,13 +145,13 @@ class JittorMPIDriver(JittorDriver):
 
     def is_distributed(self):
         """
-        判断是否为分布式的 **Driver** ，在 ``JittorSingleDriver`` 中，返回 ``True``。
+        判断是否为分布式的 **Driver** ，在 ``JittorMPIDriver`` 中，返回 ``True``。
         """
         return True
 
     @property
     def data_device(self) -> str:
         """
-        :return: 数据所在的设备；
+        :return: 数据所在的设备。
         """
         return self.model_device

fastNLP/core/drivers/jittor_driver/single_device.py

@@ -25,16 +25,21 @@ class JittorSingleDriver(JittorDriver):
     r"""
     ``Jittor`` 框架下用于 ``cpu`` 和单卡 ``gpu`` 运算的 ``Driver``。
 
-    :param model: 传入给 ``Trainer`` 的 ``model`` 参数；
-    :param device: 训练和模型所在的设备，在 **Jittor** 中，应当为以下值之一：``[None, 'cpu', 'gpu', 'cuda']``；
+    :param model: 传入给 ``Trainer`` 的 ``model`` 参数。
+    :param device: 训练和模型所在的设备，在 **Jittor** 中，应当为以下值之一：``[None, 'cpu', 'gpu', 'cuda']``：
         
-        * 为 ``None`` 或 ``cpu`` 时
-         表示在 ``cpu`` 上进行训练；
-        * 为 ``gpu`` 或 ``cuda`` 时
-         表示在显卡设备上进行训练；
+        * 为 ``None`` 或 ``cpu`` 时，表示在 ``cpu`` 上进行训练；
+        * 为 ``gpu`` 或 ``cuda`` 时，表示在显卡设备上进行训练；
 
-    :param fp16: 是否开启 fp16；
+    :param fp16: 是否开启 fp16 混合精度训练。
     :param jittor_kwargs:
+    :kwargs:
+        * *model_wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为
+
+        .. note::
+
+            关于该参数的详细说明，请参见 :class:`~fastNLP.core.controllers.Trainer` 中的描述；函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`。
+
     """
 
     def __init__(self, model, device=None, fp16: bool = False, jittor_kwargs: Dict = None, **kwargs):
@@ -50,7 +55,7 @@ class JittorSingleDriver(JittorDriver):
 
     def setup(self):
         r"""
-        初始化训练环境；根据传入的 ``device`` 值设置模型的训练场景为 ``cpu`` 或 ``gpu``；
+        初始化训练环境；根据传入的 ``device`` 值设置模型的训练场景为 ``cpu`` 或 ``gpu``。
         """
         if self.model_device in ["cpu", None]:
             jt.flags.use_cuda = 0   # 使用 cpu
@@ -130,16 +135,3 @@ class JittorSingleDriver(JittorDriver):
             return replace_batch_sampler(dataloader, batch_sampler)
         else:
             return dataloader
-
-    def unwrap_model(self):
-        """
-        返回训练使用的模型。
-        """
-        return self.model
-
-    @property
-    def data_device(self) -> str:
-        """
-        :return: 数据和模型所在的设备；
-        """
-        return self.model_device

fastNLP/core/drivers/jittor_driver/utils.py

@@ -31,7 +31,7 @@ def jittor_seed_everything(seed: int = None, add_global_rank_to_seed: bool = Tru
 
     :param seed: 全局随机状态的整数值种子。如果为 ``None`` 则会根据时间戳生成一个种子。
     :param add_global_rank_to_seed: 在分布式训练中，是否在不同 **rank** 中使用不同的随机数。
-        当设置为 ``True`` 时，**FastNLP** 会将种子加上当前的 ``global_rank``。
+        当设置为 ``True`` 时，**fastNLP** 会将种子加上当前的 ``global_rank``。
     """
     max_seed_value = np.iinfo(np.uint32).max
     min_seed_value = np.iinfo(np.uint32).min

fastNLP/core/drivers/oneflow_driver/__init__.py

@@ -1,15 +1,14 @@
 __all__ = [
-    "OneflowDDPDriver",
-    "OneflowSingleDriver",
     "OneflowDriver",
+    "OneflowSingleDriver",
+    "OneflowDDPDriver",
     "oneflow_seed_everything",
-    "optimizer_state_to_device"
 ]
 
 from .ddp import OneflowDDPDriver
 from .single_device import OneflowSingleDriver
 from .oneflow_driver import OneflowDriver
-from .utils import oneflow_seed_everything, optimizer_state_to_device
+from .utils import oneflow_seed_everything
 
 
 

fastNLP/core/drivers/oneflow_driver/ddp.py

@@ -35,7 +35,7 @@ class OneflowDDPDriver(OneflowDriver):
 
     .. note::
 
-        您在绝大多数情况下不需要自己使用到该类，通过向 ``Trainer`` 传入正确的参数，您可以方便快速地部署您的分布式训练；
+        您在绝大多数情况下不需要自己使用到该类，通过向 ``Trainer`` 传入正确的参数，您可以方便快速地部署您的分布式训练。
 
         ``OneflowDDPDriver`` 目前支持两种启动方式：
         
@@ -43,13 +43,20 @@ class OneflowDDPDriver(OneflowDriver):
             2. 用户将模型通过 ``DistributedDataParallel`` 处理后，通过运行 ``python -m oneflow.distributed.launch --nproc_per_node 2 train.py`` 启动；
 
         注意多机的启动强制要求用户在每一台机器上使用 ``python -m oneflow.distributed.launch`` 启动；因此我们不会在 ``OneflowDDPDriver`` 中保存
-        任何当前有多少台机器的信息；
+        任何当前有多少台机器的信息。
 
-    :param model: 传入给 ``Trainer`` 的 ``model`` 参数；
-    :param parallel_device: 该参数无效，**fastNLP** 会自动获取当前进程的设备；
-    :param fp16: 是否开启 fp16 训练；目前该参数无效；
+    :param model: 传入给 ``Trainer`` 的 ``model`` 参数
+    :param parallel_device: 该参数无效，**fastNLP** 会自动获取当前进程的设备
+    :param fp16: 是否开启 fp16 训练；目前该参数无效
     :param oneflow_kwargs: 
-        * *ddp_kwargs* -- 用于 ``DistributedDataParallel`` 的其它参数，详情可查阅 **oneflow** 的官方文档；
+        * *ddp_kwargs* -- 用于 ``DistributedDataParallel`` 的其它参数，详情可查阅 **oneflow** 的官方文档
+    :kwargs:
+        * *model_wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为
+
+        .. note::
+
+            关于该参数的详细说明，请参见 :class:`~fastNLP.core.controllers.Trainer` 中的描述；函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`。
+
     """
 
     def __init__(
@@ -88,7 +95,7 @@ class OneflowDDPDriver(OneflowDriver):
 
     def setup(self):
         r"""
-        将模型用 ``DistributedDataParallel`` 进行处理；
+        将模型用 ``DistributedDataParallel`` 进行处理。
         """
         if self._has_setup:
             return
@@ -121,14 +128,23 @@ class OneflowDDPDriver(OneflowDriver):
 
     @property
     def master_address(self) -> str:
+        """
+        分布式训练中的地址 ``MASTER_ADDR``
+        """
         return os.environ.get("MASTER_ADDR")
 
     @property
     def master_port(self) -> str:
+        """
+        分布式训练使用的端口 ``MASTER_PORT``
+        """
         return os.environ.get("MASTER_PORT")
 
     @property
     def world_size(self) -> int:
+        """
+        分布式训练的进程总数 ``WORLD_SIZE``
+        """
         return self._world_size
 
     @world_size.setter
@@ -137,6 +153,9 @@ class OneflowDDPDriver(OneflowDriver):
 
     @property
     def global_rank(self) -> int:
+        """
+        当前进程的全局编号 ``global_rank``
+        """
         return self._global_rank
 
     @global_rank.setter
@@ -144,11 +163,18 @@ class OneflowDDPDriver(OneflowDriver):
         self._global_rank = rank
 
     @property
-    def local_rank(self) -> int:  # 这个不会受到 all_rank_call_context 的影响
+    def local_rank(self) -> int:
+        """
+        当前进程的局部编号 ``local_rank``
+        """
         return int(os.environ.get("LOCAL_RANK", 0))
 
     @property
     def data_device(self):
+        """
+        数据所在的设备。由于 **oneflow** 可以通过 :func:`oneflow.cuda.current_device` 获取当前进程的设备，因此
+        该属性和 ``model_device`` 表现相同。
+        """
         return self._data_device
 
     def set_dist_repro_dataloader(self, dataloader,
@@ -240,13 +266,13 @@ class OneflowDDPDriver(OneflowDriver):
 
     def is_global_zero(self):
         r"""
-        :return: 返回当前的进程是否在全局上是进程 0 ；
+        :return: 当前的进程是否在全局上是进程 0 
         """
         return self.global_rank == 0
 
     def get_model_no_sync_context(self):
         r"""
-        :return: 返回一个 ``context`` 上下文环境，用于关闭各个进程之间的同步；该功能暂时无效，返回一个空的上下文环境；
+        :return: 一个 ``context`` 上下文环境，用于关闭各个进程之间的同步；该功能暂时无效，返回一个空的上下文环境
         """
         # TODO 暂时没有在 oneflow 中找到类似的功能；
         from fastNLP.core.utils import nullcontext
@@ -255,68 +281,64 @@ class OneflowDDPDriver(OneflowDriver):
 
     def unwrap_model(self):
         r"""
-        :return: 返回原始模型；
+        :return: 使用的原始模型
         """
         return self.model
 
     def get_local_rank(self) -> int:
         r"""
-        :return: 返回当前进程局部的进程编号；
+        :return: 当前进程局部的进程编号
         """
         return self.local_rank
 
     def barrier(self):
         r"""
-        通过使用该函数来使得各个进程之间同步操作；
+        同步各个进程之间的操作
         """
         if int(os.environ.get(FASTNLP_NO_SYNC, 0)) < 1:  # 当 FASTNLP_NO_SYNC 小于 1 时实际执行
             comm.barrier()
 
     def is_distributed(self):
         r"""
-        :return: 返回当前使用的 driver 是否是分布式的 driver，对于 ``OneflowDDPDriver`` 来说，该函数一定返回 ``True``；
+        :return: 当前使用的 driver 是否是分布式的 driver，对于 ``OneflowDDPDriver`` 来说，该函数一定返回 ``True``
         """
         return True
 
-    def broadcast_object(self, obj, src: int = 0, **kwargs):
+    def broadcast_object(self, obj, src: int = 0, group=None, **kwargs):
         r"""
-        从 src 端将 obj 对象（可能是 tensor ，可能是 object ）发送到 dst 处。如果是非 tensor 的对象会尝试使用 pickle 进行打包进行
-        传输，然后再 dst 处再加载回来。仅在分布式的 driver 中有实际意义。
+        从 ``src`` 端将 ``obj`` 对象（可能是 tensor ，可能是 object ）广播到其它进程。如果是非 tensor 的对象会尝试使用 pickle 进行打包进行
+        传输，然后在接收处处再加载回来。仅在分布式的 driver 中有实际意义。
 
         :param obj: obj，可能是 Tensor 或 嵌套类型的数据
-        :param int src: source 的 global rank 。
-        :param int dst: target 的 global rank，可以是多个目标 rank
-        :param group: 所属的 group
-        :return: 如果当前不是分布式 driver 直接返回输入的 obj 。如果当前 rank 是接收端（其 global rank 包含在了 dst 中），则返回
-            接收到的参数；如果是 source 端则返回发射的内容；既不是发送端、又不是接收端，则返回 None 。
+        :param src: 发送方的 ``global_rank``
+        :param group: 该参数无效
+        :return: 如果当前 rank 是接收端，则返回接收到的参数；如果是 source 端则返回发送的内容。如果环境变量 ``FASTNLP_NO_SYNC`` 为 **2** 则
+            返回 ``None``
         """
         if int(os.environ.get(FASTNLP_NO_SYNC, 0)) == 2:  # 如果 FASTNLP_NO_SYNC == 2 直接返回。
             return
         return fastnlp_oneflow_broadcast_object(obj, src, device=self.data_device)
 
-    def all_gather(self, obj) -> List:
+    def all_gather(self, obj, group) -> List:
         r"""
-        将 obj 互相传送到其它所有的 rank 上，其中 obj 可能是 Tensor，也可能是嵌套结构的 object 。如果不是基础类型的数据，尝试通过
+        将 ``obj`` 互相传送到其它所有的 rank 上，其中 ``obj`` 可能是 Tensor，也可能是嵌套结构的 object 。如果不是基础类型的数据，将会尝试通过
         pickle 进行序列化，接收到之后再反序列化。
 
         example::
 
-            obj = {
-                'a': [1, 1],
-                'b': [[1, 2], [1, 2]],
-                'c': {
-                    'd': [1, 2]
-                }
-            }
-            ->
-            [
-                {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
-                {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
-            ]
-
-        :param obj: 需要传输的对象，在每个rank上都应该保持相同的结构。
-        :param group:
-        :return:
+            >>> # rank 0
+            >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 1}}
+            >>> # rank 1
+            >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+            >>> # after all_gather():
+            >>> result = [
+                    {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
+                    {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+                ]
+
+        :param obj: 需要传输的对象，在每个 rank 上都应该保持相同的结构。
+        :param group: 该参数无效。
+        :return: 所有 rank 发送的 ``obj`` 聚合在一起的内容；如果环境变量 ``FASTNLP_NO_SYNC`` 为 **2** 则不会执行，直接返回 ``[obj]`` 。
         """
         if int(os.environ.get(FASTNLP_NO_SYNC, 0)) == 2:  # 如果 FASTNLP_NO_SYNC 表示不执行
             return [obj]

fastNLP/core/drivers/oneflow_driver/dist_utils.py

@@ -13,6 +13,8 @@ if _NEED_IMPORT_ONEFLOW:
 
 PROTOCOL_VERSION = 1
 
+__all__ = []
+
 def _validate_output_list_for_rank(my_rank, dst, gather_list):
     if dst == my_rank:
         if not gather_list:
@@ -176,18 +178,15 @@ def fastnlp_oneflow_all_gather(obj: Any, device=None) ->List:
 
     example::
 
-        obj = {
-            'a': [1, 1],
-            'b': [[1, 2], [1, 2]],
-            'c': {
-                'd': [1, 2]
-            }
-        }
-        ->
-        [
-            {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
-            {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
-        ]
+        >>> # rank 0
+        >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 1}}
+        >>> # rank 1
+        >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+        >>> # after all_gather():
+        >>> result = [
+                {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
+                {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+            ]
 
     :param obj: 任意结构的数据，如果为 tensor ，需要保证每个显卡上的 tensor 的形状是一样的。如果传入的是非 tensor 对象都将直接进行
         序列化之后进行传输。

fastNLP/core/drivers/oneflow_driver/initialize_oneflow_driver.py

@@ -18,11 +18,11 @@ def initialize_oneflow_driver(driver: str, device: Optional[Union[str, "oneflow.
     r"""
     用来根据参数 ``driver` 和 ``device`` 来确定并且初始化一个具体的 ``Driver`` 实例然后返回回去；
 
-    :param driver: 该参数的值应为以下之一：``["oneflow"]``；
-    :param device: 该参数的格式与 ``Trainer`` 对参数 ``device`` 的要求一致；
+    :param driver: 该参数的值应为以下之一：``["oneflow"]``
+    :param device: 该参数的格式与 ``Trainer`` 对参数 ``device`` 的要求一致
     :param model: 训练或者评测的具体的模型；
 
-    :return: 返回一个 :class:`~fastNLP.core.OneflowSingleDriver` 或 :class:`~fastNLP.core.OneflowDDPDriver` 实例；
+    :return: 一个 :class:`~fastNLP.core.OneflowSingleDriver` 或 :class:`~fastNLP.core.OneflowDDPDriver` 实例；
     """
     # world_size 和 rank
     if FASTNLP_BACKEND_LAUNCH in os.environ:

fastNLP/core/drivers/oneflow_driver/oneflow_driver.py

@@ -36,17 +36,23 @@ from fastNLP.core.dataloaders import OverfitDataLoader
 
 class OneflowDriver(Driver):
     r"""
-    专属于 ``oneflow`` 的 ``driver``，是 ``OneflowSingleDriver`` 和 ``OneflowDDPDriver`` 的父类；
+    实现了 **oneflow** 框架训练功能的基本 ``Driver``。这个类被以下子类继承：
+    
+        1. :class:`~fastNLP.core.drivers.oneflow_driver.OneflowSingleDriver` ：实现了使用单卡和 ``cpu`` 训练的具体功能；
+        2. :class:`~fastNLP.core.drivers.oneflow_driver.OneflowDDPDriver` ：实现了使用 ``DistributedDataParallel`` 启动 **oneflow** 分布式训练的功能；
 
     .. warning::
 
         您不应当直接初始化该类，然后传入给 ``Trainer``，换句话说，您应当使用该类的子类 ``OneflowSingleDriver`` 和 ``OneflowDDPDriver``，而不是
-        该类本身；
+        该类本身。
 
     .. note::
 
-        您可以在使用 ``OneflowSingleDriver`` 和 ``OneflowDDPDriver`` 时使用 ``OneflowDriver`` 提供的接口；
+        您可以在使用 ``OneflowSingleDriver`` 和 ``OneflowDDPDriver`` 时使用 ``OneflowDriver`` 提供的接口。
 
+    :param model: 训练使用的模型
+    :param fp16: 该参数暂时无效
+    :param oneflow_kwargs:
     """
     def __init__(self, model, fp16: Optional[bool] = False, oneflow_kwargs: Dict = None, **kwargs):
         super(OneflowDriver, self).__init__(model)
@@ -66,19 +72,33 @@ class OneflowDriver(Driver):
         self.wo_auto_param_call = kwargs.get("model_wo_auto_param_call", False)
 
     def zero_grad(self):
+        """
+        实现梯度置零的过程
+        """
         for optimizer in self.optimizers:
             optimizer.zero_grad(self.set_grad_to_none)
 
     def backward(self, loss):
+        """
+        对 ``loss`` 进行反向传播
+        """
         loss.backward()
         # self.grad_scaler.scale(loss).backward()
 
     def step(self):
+        r"""
+        实现参数的优化更新过程
+        """
         for optimizer in self.optimizers:
             self.grad_scaler.step(optimizer)
             self.grad_scaler.update()
 
     def check_dataloader_legality(self, dataloader):
+        """
+        检测 DataLoader 是否合法。支持的类型包括 :class:`~fastNLP.core.dataloaders.OneflowDataLoader`、 :class:`oneflow.utils.data.DataLoader` 。
+
+        :param dataloder:
+        """
         if not isinstance(dataloader, DataLoader) and not isinstance(dataloader, OverfitDataLoader):
             raise TypeError(f"{DataLoader} is expected, instead of `{type(dataloader)}`")
         if len(dataloader) == 0:
@@ -95,11 +115,11 @@ class OneflowDriver(Driver):
     @staticmethod
     def tensor_to_numeric(tensor, reduce: str = None):
         r"""
-        将 ``oneflow.Tensor`` 转换成 python 中的数值类型；
+        将 ``oneflow.Tensor`` 转换成 python 中的数值类型。
 
-        :param tensor: ``oneflow.Tensor``；
-        :param reduce: 当 tensor 是一个多数值的张量时，应当使用何种归一化操作来转换成单一数值，应当为以下类型之一：``['max', 'min', 'sum', 'mean']``；
-        :return: 返回一个单一数值，其数值类型是 python 中的基本的数值类型，例如 ``int，float`` 等；
+        :param tensor: ``oneflow.Tensor``
+        :param reduce: 当 tensor 是一个多数值的张量时，应当使用何种归一化操作来转换成单一数值，应当为以下类型之一：``['max', 'min', 'sum', 'mean']``
+        :return: 一个单一数值，其数值类型是 python 中的基本的数值类型，例如 ``int，float`` 等。
         """
 
         if tensor is None:
@@ -120,8 +140,9 @@ class OneflowDriver(Driver):
 
     def set_model_mode(self, mode: str):
         r"""
-        设置模型的状态是 ``train`` 还是 ``eval``；
-        :param mode: ``'train'`` 或 ``'eval'``；
+        设置模型为 ``train`` 或 ``eval`` 的模式；目的是为切换模型的训练和推理（会关闭 dropout 等）模式。
+
+        :param mode: 应为二者之一：``["train", "eval"]``
         """
         assert mode in {"train", "eval"}
         getattr(self.model, mode)()
@@ -129,10 +150,10 @@ class OneflowDriver(Driver):
     @rank_zero_call
     def save_model(self, filepath: Union[str, Path], only_state_dict: bool = True, **kwargs):
         """
-        保存当前 driver 的模型到 folder 下。
+        保存当前 driver 的模型到 ``filepath``。
 
-        :param filepath: 保存到哪个文件夹；
-        :param only_state_dict: 是否只保存权重；如果使用 ``DistributedDataParallel`` 启动分布式训练的话，该参数只能为 ``True``；
+        :param filepath: 保存文件的文件位置
+        :param only_state_dict: 是否只保存权重；如果使用 ``DistributedDataParallel`` 启动分布式训练的话，该参数只能为 ``True``
         :return:
         """
         model = self.unwrap_model()
@@ -155,12 +176,11 @@ class OneflowDriver(Driver):
 
     def load_model(self, filepath: Union[Path, str], only_state_dict: bool = True, **kwargs):
         """
-        从 folder 中加载权重并赋值到当前 driver 的模型上。
+        加载模型的函数；将 ``filepath`` 中的模型加载并赋值给当前 ``model`` 。
 
-        :param filepath: 加载权重或模型的路径
-        :param load_state_dict: 保存的内容是否只是权重。
-        :param kwargs:
-        :return:
+        :param filepath: 保存文件的文件位置
+        :param load_state_dict: 保存的内容是否只是权重；如果使用 ``DistributedDataParallel`` 启动分布式训练的话，
+            该参数只能为 ``True``
         """
         model = self.unwrap_model()
         res = oneflow.load(filepath)
@@ -176,6 +196,17 @@ class OneflowDriver(Driver):
 
     @rank_zero_call
     def save_checkpoint(self, folder: Path, states: Dict, dataloader, only_state_dict: bool = True, should_save_model: bool = True, **kwargs):
+        r"""
+        断点重训的保存函数，该函数会负责保存 **优化器** 和 **sampler** 的状态，以及 **模型** （若 ``should_save_model`` 为 ``True``）
+
+        :param folder: 保存断点重训的状态的文件夹；:meth:`save_checkpoint` 函数应该在该路径下面下面新增名为 ``FASTNLP_CHECKPOINT_FILENAME`` 与
+            ``FASTNLP_MODEL_FILENAME`` （如果 ``should_save_model`` 为 ``True`` ）的文件。把 model 相关的内容放入到 ``FASTNLP_MODEL_FILENAME`` 文件
+            中，将传入的 ``states`` 以及自身产生的其它状态一并保存在 ``FASTNLP_CHECKPOINT_FILENAME`` 里面。
+        :param states: 由 :class:`~fastNLP.core.controllers.Trainer` 传入的一个字典，其中已经包含了为了实现断点重训所需要保存的其它对象的状态。
+        :param dataloader: 正在使用的 dataloader。
+        :param only_state_dict: 是否只保存模型的参数，当 ``should_save_model`` 为 ``False`` ，该参数无效。
+        :param should_save_model: 是否应该保存模型，如果为 ``False`` ，Driver 将不负责 model 的保存。
+        """
         # 传入的 dataloader 参数是 trainer 的 dataloader 属性，因为 driver 的所有 dataloader 我们是不会去改变它的，而是通过改变
         #  trainer.dataloader 来改变 dataloader 的状态，从而适配训练或者评测环境；
 
@@ -280,6 +311,36 @@ class OneflowDriver(Driver):
         logger.debug("Load optimizer state dict.")
 
     def load_checkpoint(self, folder: Path, dataloader, only_state_dict: bool = True, should_load_model: bool = True, **kwargs) -> Dict:
+        r"""
+        断点重训的加载函数，该函数会负责读取数据，并且恢复 **优化器** 、**sampler** 的状态和 **模型** （如果 ``should_load_model`` 为 True）以及其它
+        在 :meth:`save_checkpoint` 函数中执行的保存操作，然后将一个 state 字典返回给 :class:`~fastNLP.core.controllers.Trainer` （ 内容为 :meth:`save_checkpoint` 
+        接受到的 ``states`` ）。
+
+        该函数应该在所有 rank 上执行。
+
+        :param folder: 读取该 folder 下的 ``FASTNLP_CHECKPOINT_FILENAME`` 文件与 ``FASTNLP_MODEL_FILENAME``
+            （如果 should_load_model 为True）。
+        :param dataloader: 当前给定 dataloader，需要根据保存的 dataloader 状态合理设置。若该值为 ``None`` ，则不需要返回 ``'dataloader'``
+            以及 ``'batch_idx_in_epoch'`` 这两个值。
+        :param only_state_dict: 是否仅读取模型的 state_dict ，当 ``should_save_model`` 为 ``False`` ，该参数无效。如果为 ``True`` ，说明保存的内容为权重；如果为
+            False 说明保存的是模型，但也是通过当前 Driver 的模型去加载保存的模型的权重，而不是使用保存的模型替换当前模型。
+        :param should_load_model: 是否应该加载模型，如果为 ``False`` ，Driver 将不负责加载模型。若该参数为 ``True`` ，但在保存的状态中没有
+            找到对应的模型状态，则报错。
+        :return: :meth:`save_checkpoint` 函数输入的 ``states`` 内容。除此之外，还返回的内容有：
+
+            * *dataloader* -- 根据传入的 ``dataloader`` 与读取出的状态设置为合理状态的 dataloader。在当前 ``dataloader`` 样本数与读取出的 sampler 样本数
+              不一致时报错。
+            * *batch_idx_in_epoch* -- :class:`int` 类型的数据，表明当前 epoch 进行到了第几个 batch 。请注意，该值不能仅通过保存的数据中读取的，因为前后两次运行的
+              ``batch_size`` 可能有变化，而应该符合以下等式::
+
+                返回的 dataloader 还会产生的 batch 数量 + batch_idx_in_epoch = 原来不断点训练时的 batch 的总数
+              
+              由于 ``返回的 dataloader 还会产生的batch数`` 在 ``batch_size`` 与 ``drop_last`` 参数给定的情况下，无法改变，因此只能通过调整 ``batch_idx_in_epoch``
+              这个值来使等式成立。一个简单的计算原则如下：
+
+                * drop_last 为 ``True`` 时，等同于 floor(sample_in_this_rank/batch_size) - floor(num_left_samples/batch_size)；
+                * drop_last 为 ``False`` 时，等同于 ceil(sample_in_this_rank/batch_size) - ceil(num_left_samples/batch_size)。
+        """
         states = oneflow.load(folder.joinpath(FASTNLP_CHECKPOINT_FILENAME))
 
         # 1. 加载 optimizers 的状态；
@@ -309,7 +370,9 @@ class OneflowDriver(Driver):
 
     def get_evaluate_context(self):
         r"""
-        :return: 返回 ``oneflow.no_grad`` 这个 context；
+        返回一个不计算梯度的上下文环境用来对模型进行评测。
+
+        :return: 上下文对象 ``oneflow.no_grad``
         """
         return oneflow.no_grad
 
@@ -335,17 +398,17 @@ class OneflowDriver(Driver):
     @staticmethod
     def move_model_to_device(model: "oneflow.nn.Module", device: "oneflow.device"):
         r"""
-        将模型迁移到对应的设备上；
+        将模型迁移到对应的设备上。
         """
         if device is not None:
             model.to(device)
 
     def move_data_to_device(self, batch):
         """
-        将一个 batch 的数据迁移到对应的设备上；
+        将一个 ``batch`` 的数据迁移到对应的设备上。
 
-        :param batch: 一个 batch 的数据，可以是 ``list、dict`` 等；
-        :return:
+        :param batch: 包含 :class:`oneflow.Tensor` 的数据集合，可以是 **List**、**Dict** 等嵌套类型
+        :return: 移动到指定机器后的 ``batch``
         """
         return oneflow_move_data_to_device(batch, self.data_device)
 
@@ -366,11 +429,20 @@ class OneflowDriver(Driver):
         random.seed(stdlib_seed)
 
     def set_deterministic_dataloader(self, dataloader: "DataLoader"):
+        """
+        为了确定性训练要对 ``dataloader`` 进行修改，保证在确定随机数种子后，每次重新训练得到的结果是一样的。 
+        """
         if dataloader.worker_init_fn is None:
             dataloader.worker_init_fn = partial(self.worker_init_function,
                                                 rank=int(os.environ.get(FASTNLP_GLOBAL_RANK, 0)))
 
     def set_sampler_epoch(self, dataloader: "DataLoader", cur_epoch_idx: int):
+        r"""
+        对于分布式的 ``sampler``，需要在每一个 ``epoch`` 前设置随机数种子，来保证每一个进程上的 ``shuffle`` 是一样的。
+
+        :param dataloader: 需要设置 ``epoch`` 的 ``dataloader``
+        :param cur_epoch_idx: 当前是第几个 ``epoch``
+        """
         # 保证 ddp 训练时的 shuffle=True 时的正确性，因为需要保证每一个进程上的 sampler 的shuffle 的随机数种子是一样的；
         if callable(getattr(dataloader.sampler, "set_epoch", None)):
             dataloader.sampler.set_epoch(cur_epoch_idx)
@@ -378,9 +450,9 @@ class OneflowDriver(Driver):
     @staticmethod
     def get_dataloader_args(dataloader: "DataLoader"):
         """
-        获取 dataloader 的 shuffle 和 drop_last 属性；
+        从 ``dataloader`` 中获取参数 ``dataset``, ``batch_sampler``, ``sampler``, ``batch_size``, ``shuffle`` 
+        和 ``drop_last`` 。
         """
-
         @dataclass
         class Res:
             dataset: Optional[Dataset] = None

fastNLP/core/drivers/oneflow_driver/single_device.py

@@ -21,12 +21,19 @@ from fastNLP.core.log import logger
 
 class OneflowSingleDriver(OneflowDriver):
     r"""
-    用于执行 ``oneflow`` 动态图 cpu 和 单卡 gpu 运算的 ``driver``；
+    用于执行 ``oneflow`` 动态图 cpu 和 单卡 gpu 运算的 ``driver``。
 
-    :param model: 传入给 ``Trainer`` 的 ``model`` 参数；
-    :param device: oneflow.device，当前进程所使用的设备；
-    :param fp16: 是否开启 fp16；目前动态图的单卡下该参数无效；
+    :param model: 传入给 ``Trainer`` 的 ``model`` 参数
+    :param device: oneflow.device，当前进程所使用的设备
+    :param fp16: 是否开启 fp16；目前动态图的单卡下该参数无效。
     :param oneflow_kwargs:
+    :kwargs:
+        * *model_wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为。
+
+        .. note::
+
+            关于该参数的详细说明，请参见 :class:`~fastNLP.core.controllers.Trainer` 中的描述；函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`。
+
     """
 
     def __init__(self, model, device: "oneflow.device", fp16: bool = False, oneflow_kwargs: Dict = None, **kwargs):
@@ -54,7 +61,7 @@ class OneflowSingleDriver(OneflowDriver):
 
     def setup(self):
         r"""
-        将模型迁移到相应的设备上；
+        将模型迁移到相应的设备上。
         """
         if self.model_device is not None:
             self.model.to(self.model_device)
@@ -96,19 +103,19 @@ class OneflowSingleDriver(OneflowDriver):
 
     def unwrap_model(self):
         r"""
-        :return: 返回模型
+        :return: 训练使用的模型
         """
         return self.model
 
     @property
     def data_device(self):
         r"""
-        :return: 数据和模型所在的设备；
+        :return: 数据和模型所在的设备。
         """
         return self.model_device
 
     def is_distributed(self):
         r"""
-        :return: 返回当前使用的 driver 是否是分布式的 driver，在 ``OneflowSingleDriver`` 中返回 ``False``；
+        :return: 当前使用的 driver 是否是分布式的 driver，在 ``OneflowSingleDriver`` 中返回 ``False``。
         """
         return False

fastNLP/core/drivers/oneflow_driver/utils.py

@@ -31,7 +31,6 @@ else:
 
 __all__ = [
     'oneflow_seed_everything',
-    'optimizer_state_to_device'
 ]
 
 def oneflow_seed_everything(seed: int = None, add_global_rank_to_seed: bool = True) -> int:
@@ -40,7 +39,7 @@ def oneflow_seed_everything(seed: int = None, add_global_rank_to_seed: bool = Tr
 
     :param seed: 全局随机状态的整数值种子。如果为 ``None`` 则会根据时间戳生成一个种子。
     :param add_global_rank_to_seed: 在分布式训练中，是否在不同 **rank** 中使用不同的随机数。
-        当设置为 ``True`` 时，**FastNLP** 会将种子加上当前的 ``global_rank``。
+        当设置为 ``True`` 时，**fastNLP** 会将种子加上当前的 ``global_rank``。
     """
     max_seed_value = np.iinfo(np.uint32).max
     min_seed_value = np.iinfo(np.uint32).min
@@ -263,11 +262,11 @@ def replace_batch_sampler(dataloader, new_batch_sampler):
 
 def optimizer_state_to_device(state, device):
     r"""
-    将一个 ``optimizer`` 的 ``state_dict`` 迁移到对应的设备；
+    将一个 ``optimizer`` 的 ``state_dict`` 迁移到对应的设备。
 
-    :param state: ``optimzier.state_dict()``；
-    :param device: 要迁移到的目的设备；
-    :return: 返回迁移后的新的 state_dict；
+    :param state: :func:`optimzier.state_dict` 获取的 state_dictt
+    :param device: 要迁移到的目的设备。
+    :return: 迁移后的新的 state_dict。
     """
     new_state = {}
     for name, param in state.items():

fastNLP/core/drivers/paddle_driver/dist_utils.py

@@ -175,18 +175,15 @@ def fastnlp_paddle_all_gather(obj: Any, device=None, group=None) ->List:
 
     example::
 
-        obj = {
-            'a': [1, 1],
-            'b': [[1, 2], [1, 2]],
-            'c': {
-                'd': [1, 2]
-            }
-        }
-        ->
-        [
-            {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
-            {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
-        ]
+        >>> # rank 0
+        >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 1}}
+        >>> # rank 1
+        >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+        >>> # after all_gather():
+        >>> result = [
+                {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
+                {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+            ]
 
     :param obj: 任意结构的数据，如果为 tensor ，需要保证每个显卡上的 tensor 的形状是一样的。如果传入的是非 tensor 对象都将直接进行
         序列化之后进行传输。

fastNLP/core/drivers/paddle_driver/fleet.py

@@ -61,7 +61,7 @@ r"""
 .. note::
 
     多机的启动强制要求用户在每一台机器上使用 ``python -m paddle.distributed.launch`` 启动；因此我们不会在 ``PaddleFleetDriver``
-    中保存任何当前有多少台机器的信息；
+    中保存任何当前有多少台机器的信息。
 
 """
 import os
@@ -82,7 +82,6 @@ from fastNLP.core.utils import (
     auto_param_call,
     check_user_specific_params,
     is_in_paddle_dist,
-    is_in_paddle_dist,
     get_paddle_device_id,
 )
 from fastNLP.core.utils.paddle_utils import _convert_data_device
@@ -120,26 +119,26 @@ __all__ = [
 
 class PaddleFleetDriver(PaddleDriver):
     """
-    :param model: 训练使用的模型；
+    :param model: 训练使用的模型。
 
         * 如果不想自己初始化分布式环境，类型应为 :class:`paddle.nn.Layer`；
         * 如果已经在外面初始化了分布式环境，类型应为 :class:`paddle.DataParallel`；
 
     :param parallel_device: 多卡训练时使用的设备，必须是一个列表。
-        当使用 ``python -m paddle.distributed.launch`` 启动时，该参数无效；
+        当使用 ``python -m paddle.distributed.launch`` 启动时，该参数无效。
     :param is_pull_by_paddle_run: 标记当前进程是否为通过 ``python -m paddle.distributed.launch`` 启动的。
         这个参数仅在 :class:`~fastNLP.core.Trainer` 中初始化 driver 时使用
-    :param fp16: 是否开启混合精度训练；
+    :param fp16: 是否开启混合精度训练
     :param paddle_kwargs:
         * *fleet_kwargs* -- 用于在使用 ``PaddleFleetDriver`` 时指定 ``DataParallel`` 和 ``fleet`` 初始化时的参数，包括：
             
-            * *is_collective* -- 是否使用 paddle 集群式的分布式训练方法，目前仅支持为 ``True`` 的情况；
-            * *role_maker* -- 初始化 ``fleet`` 分布式训练 API 时使用的 ``RoleMaker``；
-            * 其它用于初始化 ``DataParallel`` 的参数；
-        * *gradscaler_kwargs* -- 用于 ``fp16=True`` 时，提供给 :class:`paddle.amp.GradScaler` 的参数;
+            * *is_collective* -- 是否使用 paddle 集群式的分布式训练方法，目前仅支持为 ``True`` 的情况。
+            * *role_maker* -- 初始化 ``fleet`` 分布式训练 API 时使用的 ``RoleMaker``。
+            * 其它用于初始化 ``DataParallel`` 的参数。
+        * *gradscaler_kwargs* -- 用于 ``fp16=True`` 时，提供给 :class:`paddle.amp.GradScaler` 的参数
     
     :kwargs:
-        * wo_auto_param_call (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为；
+        * *wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为
 
         .. note::
 
@@ -176,9 +175,9 @@ class PaddleFleetDriver(PaddleDriver):
         self.parallel_device = parallel_device
         # 在初始化时，如果发现 is_pull_by_paddle_run ，则将 parallel_device 设置成当前进程的gpu
         if is_pull_by_paddle_run:
-            self._model_device = parallel_device
+            self.model_device = parallel_device
         else:
-            self._model_device = parallel_device[self.local_rank]
+            self.model_device = parallel_device[self.local_rank]
 
         # 如果用户自己在外面初始化了并行模型；
         self.outside_fleet = False
@@ -311,6 +310,9 @@ class PaddleFleetDriver(PaddleDriver):
         self.global_rank = paddledist.get_rank()
 
     def barrier(self):
+        """
+        同步进程之间的操作
+        """
         if int(os.environ.get(FASTNLP_NO_SYNC, 0)) < 1:  # 当 FASTNLP_NO_SYNC 小于 1 时实际执行
             paddledist.barrier()
 
@@ -329,6 +331,9 @@ class PaddleFleetDriver(PaddleDriver):
 
     @property
     def world_size(self) -> int:
+        """
+        分布式训练的进程总数 ``WOLRD_SIZE``
+        """
         return self._world_size
 
     @world_size.setter
@@ -337,6 +342,9 @@ class PaddleFleetDriver(PaddleDriver):
 
     @property
     def global_rank(self) -> int:
+        """
+        当前进程的全局编号 ``global_rank``
+        """
         return self._global_rank
 
     @global_rank.setter
@@ -345,20 +353,16 @@ class PaddleFleetDriver(PaddleDriver):
 
     @property
     def local_rank(self) -> int:
-        return int(os.getenv("PADDLE_RANK_IN_NODE", "0"))
-
-    @property
-    def model_device(self):
         """
-        :return: 模型所在的设备；
+        当前进程的局部编号 ``local_rank``
         """
-        return self._model_device
+        return int(os.getenv("PADDLE_RANK_IN_NODE", "0"))
 
     @property
     def data_device(self):
         """
-        :return: 数据所在的设备；由于 **PaddlePaddle** 可以通过环境变量获取当前进程的设备，因此该属性
-            和 ``model_device`` 表现相同；
+        数据所在的设备；由于 **PaddlePaddle** 可以通过环境变量获取当前进程的设备，因此该属性
+        和 ``model_device`` 表现相同。
         """
         return self.model_device
 
@@ -484,9 +488,15 @@ class PaddleFleetDriver(PaddleDriver):
             raise ValueError("Parameter `dist_sampler` can only be one of three values: ('dist', 'unrepeatdist', None).")
 
     def is_global_zero(self) -> bool:
+        r"""
+        :return: 当前的进程是否在全局上是进程 0
+        """
         return self.global_rank == 0
 
     def get_model_no_sync_context(self):
+        r"""
+        :return: 一个 ``context`` 上下文环境，用于关闭各个进程之间的同步。
+        """
         return self.model.no_sync
 
     def unwrap_model(self) -> "paddle.nn.Layer":
@@ -500,11 +510,14 @@ class PaddleFleetDriver(PaddleDriver):
             return _layers
 
     def get_local_rank(self) -> int:
+        r"""
+        :return: 当前进程局部的进程编号。
+        """
         return self.local_rank
 
     def is_distributed(self) -> bool:
         """
-        判断是否为分布式的 **Driver** ，在 ``PaddleFleetDriver`` 中，返回 ``True``。
+        :return: 当前使用的 driver 是否是分布式的 driver，在 ``PaddleFleetDriver`` 中，返回 ``True``。
         """
         return True
 
@@ -518,9 +531,39 @@ class PaddleFleetDriver(PaddleDriver):
                                 f"not {type(each_optimizer)}.")
 
     def broadcast_object(self, obj, src:int=0, group=None, **kwargs):
+        r"""
+        从 ``src`` 端将 ``obj`` 对象（可能是 tensor ，可能是 object ）广播到其它进程。如果是非 tensor 的对象会尝试使用 pickle 进行打包进行
+        传输，然后在接收处处再加载回来。仅在分布式的 driver 中有实际意义。
+
+        :param obj: obj，可能是 Tensor 或 嵌套类型的数据
+        :param src: 发送方的 ``global_rank``
+        :param group: 进程所在的通信组
+        :return: 如果当前 rank 是接收端，则返回接收到的参数；如果是 source 端则返回发送的内容。如果环境变量 ``FASTNLP_NO_SYNC`` 为 **2** 则
+            返回 ``None``
+        """
         # 因为设置了CUDA_VISIBLE_DEVICES，可能会引起错误
         device = _convert_data_device(self.data_device)
         return fastnlp_paddle_broadcast_object(obj, src, device=device, group=group)
 
     def all_gather(self, obj, group=None) -> List:
+        r"""
+        将 ``obj`` 互相传送到其它所有的 rank 上，其中 ``obj`` 可能是 Tensor，也可能是嵌套结构的 object 。如果不是基础类型的数据，将会尝试通过
+        pickle 进行序列化，接收到之后再反序列化。
+
+        example::
+
+            >>> # rank 0
+            >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 1}}
+            >>> # rank 1
+            >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+            >>> # after all_gather():
+            >>> result = [
+                    {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
+                    {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+                ]
+
+        :param obj: 需要传输的对象，在每个 rank 上都应该保持相同的结构。
+        :param group: 进程所在的通信组。
+        :return: 所有 rank 发送的 ``obj`` 聚合在一起的内容；如果环境变量 ``FASTNLP_NO_SYNC`` 为 **2** 则不会执行，直接返回 ``[obj]`` 。
+        """
         return fastnlp_paddle_all_gather(obj, group=group)

fastNLP/core/drivers/paddle_driver/initialize_paddle_driver.py

@@ -27,8 +27,8 @@ def initialize_paddle_driver(driver: str, device: Optional[Union[str, int, List[
     2. 如果 ``device`` 包含了多个设备，则返回一个 :class:`~fastNLP.core.PaddleFleetDriver` 实例，否则返回
     单卡的 :class:`~fastNLP.core.PaddleSingleDriver` 实例
 
-    :param driver: 使用的 ``driver`` 类型，在这个函数中仅支持 ``paddle``；
-    :param device: 该参数的格式与 ``Trainer`` 对参数 ``device`` 的要求一致；
+    :param driver: 使用的 ``driver`` 类型，在这个函数中仅支持 ``paddle``
+    :param device: 该参数的格式与 ``Trainer`` 对参数 ``device`` 的要求一致
     :param model: 训练或者评测的具体的模型；
 
     :return: 一个 :class:`~fastNLP.core.PaddleSingleDriver` 或 :class:`~fastNLP.core.PaddleFleetDriver` 实例；

fastNLP/core/drivers/paddle_driver/paddle_driver.py

@@ -47,28 +47,25 @@ if _NEED_IMPORT_PADDLE:
 
 class PaddleDriver(Driver):
     r"""
-    实现了 **PaddlePaddle** 框架训练功能的基本 Driver，实现了单卡和多卡情景下均需要实现的功能，以和 **fastNLP** 的
-    :class:`~fastNLP.core.Trainer` 兼容；通过这个 Driver，可以在 **fastNLP** 中实现从 **Pytorch** 框架到
-    **PaddlePaddle** 深度学习框架的切换。
+    实现了 **PaddlePaddle** 框架训练功能的基本 Driver。
 
     这个类被以下子类继承：
 
-    1. :class:`~fastNLP.core.drivers.PaddleSingleDriver`：实现了使用单卡和 ``cpu`` 训练的具体功能；
-    2. :class:`~fastNLP.core.drivers.PaddleFleetDriver`：实现了使用 ``fleet`` 分布式训练 API 进行集群式分布式训练的具体功能；
+    1. :class:`~fastNLP.core.drivers.paddle_driver.PaddleSingleDriver`：实现了使用单卡和 ``cpu`` 训练的具体功能；
+    2. :class:`~fastNLP.core.drivers.paddle_driver.PaddleFleetDriver`：实现了使用 ``fleet`` 分布式训练 API 进行集群式分布式训练的具体功能；
 
     .. warning::
 
         您不应当直接初始化该类，然后传入给 ``Trainer``，换句话说，您应当使用该类的子类 ``PaddleSingleDriver`` 和 ``PaddleDDPDriver``，而不是
-        该类本身；
+        该类本身。
 
     .. note::
 
-        您可以在使用 ``PaddleSingleDriver`` 和 ``PaddleFleetDriver`` 时使用 ``PaddleDriver`` 提供的接口；
+        您可以在使用 ``PaddleSingleDriver`` 和 ``PaddleFleetDriver`` 时使用 ``PaddleDriver`` 提供的接口。
 
-    :param model: 训练时使用的 **PaddlePaddle** 模型；
-    :param fp16: 是否开启混合精度训练；
+    :param model: 训练时使用的 **PaddlePaddle** 模型
+    :param fp16: 是否开启混合精度训练
     :param paddle_kwargs:
-
     """
     def __init__(self, model: "paddle.nn.Layer", fp16: Optional[bool] = False, paddle_kwargs: Dict = None, **kwargs):
         if not isinstance(model, paddle.nn.Layer):
@@ -87,18 +84,32 @@ class PaddleDriver(Driver):
         self.wo_auto_param_call = kwargs.get("model_wo_auto_param_call", False)
 
     def zero_grad(self):
+        """
+        实现梯度置零的过程
+        """
         for optimizer in self.optimizers:
             optimizer.clear_grad()
 
     def backward(self, loss):
+        """
+        对 ``loss`` 进行反向传播
+        """  
         self.grad_scaler.scale(loss).backward()
 
     def step(self):
+        r"""
+        实现参数的优化更新过程
+        """
         for optimizer in self.optimizers:
             self.grad_scaler.step(optimizer)
             self.grad_scaler.update()
 
     def check_dataloader_legality(self, dataloader):
+        """
+        检测 DataLoader 是否合法。支持的类型包括 :class:`~fastNLP.core.dataloaders.PaddleDataLoader`、 :class:`paddle.io.DataLoader` 。
+
+        :param dataloder:
+        """
         if not isinstance(dataloader, DataLoader) and not isinstance(dataloader, OverfitDataLoader):
             raise TypeError(f"{DataLoader} is expected, instead of `{type(dataloader)}`")
         if dataloader.batch_size is None and dataloader.batch_sampler is None:
@@ -113,7 +124,7 @@ class PaddleDriver(Driver):
         r"""
         对于用户传入 trainer 的每一个 optimizer检测其合法性，必须为`paddle.optimizer.Optimizer`类型。
 
-        :param optimizers: 需要检测的 `optimizers`；
+        :param optimizers: 需要检测的 `optimizers`。
         """
         for each_optimizer in optimizers:
             if not isinstance(each_optimizer, Optimizer):
@@ -123,11 +134,11 @@ class PaddleDriver(Driver):
     @staticmethod
     def tensor_to_numeric(tensor, reduce=None):
         r"""
-        将一个 :class:`paddle.Tensor` 对象转换为 转换成 python 中的数值类型；
+        将一个 :class:`paddle.Tensor` 对象转换为 转换成 python 中的数值类型。
 
-        :param tensor: :class:`paddle.Tensor` 类型的对象；
-        :param reduce: 当 tensor 是一个多数值的张量时，应当使用何种归一化操作来转换成单一数值，应当为以下类型之一：``['max', 'min', 'sum', 'mean']``；
-        :return: 返回一个单一数值，其数值类型是 python 中的基本的数值类型，例如 ``int，float`` 等；
+        :param tensor: :class:`paddle.Tensor` 类型的对象。
+        :param reduce: 当 tensor 是一个多数值的张量时，应当使用何种归一化操作来转换成单一数值，应当为以下类型之一：``['max', 'min', 'sum', 'mean']``。
+        :return: 一个单一数值，其数值类型是 python 中的基本的数值类型，例如 ``int，float`` 等。
         """
         if tensor is None:
             return None
@@ -148,6 +159,11 @@ class PaddleDriver(Driver):
         )
 
     def set_model_mode(self, mode: str):
+        r"""
+        设置模型为 ``train`` 或 ``eval`` 的模式；目的是为切换模型的训练和推理（会关闭 dropout 等）模式。
+
+        :param mode: 应为二者之一：``["train", "eval"]``
+        """
         assert mode in {"train", "eval"}
         getattr(self.model, mode)()
 
@@ -156,13 +172,13 @@ class PaddleDriver(Driver):
         r"""
         将模型保存到 ``filepath`` 中。
 
-        :param filepath: 保存文件的文件位置（需要包括文件名）；
+        :param filepath: 保存文件的文件位置（需要包括文件名）。
         :param only_state_dict: 是否只保存模型的 ``state_dict``；如果为 ``False``，则会调用 ``paddle.jit.save`` 
-            函数保存整个模型的参数，此时需要传入 ``input_spec`` 参数；
+            函数保存整个模型的参数，此时需要传入 ``input_spec`` 参数。
         :kwargs:
             * *input_spec* -- 描述存储模型 ``forward`` 方法的输入；
-            当 ``only_state_dict`` 为 ``False`` 时必须传入，否则加载时会报错。您可以通过 ``InputSpec`` 或者示例 ``Tensor``
-            进行描述。详细的使用方法可以参考 **PaddlePaddle** `关于 paddle.jit.save 函数的文档 <https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/jit/save_cn.html#save>`_；
+              当 ``only_state_dict`` 为 ``False`` 时必须传入，否则加载时会报错。您可以通过 ``InputSpec`` 或者示例 ``Tensor``
+              进行描述。详细的使用方法可以参考 **PaddlePaddle** `关于 paddle.jit.save 函数的文档 <https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/jit/save_cn.html#save>`_。
         """
         model = self.unwrap_model()
         if isinstance(filepath, Path):
@@ -178,6 +194,12 @@ class PaddleDriver(Driver):
             paddle.jit.save(model, filepath, input_spec)
 
     def load_model(self, filepath: Union[Path, str], only_state_dict: bool = True, **kwargs):
+        """
+        加载模型的函数；将 ``filepath`` 中的模型加载并赋值给当前 ``model`` 。
+
+        :param filepath: 保存文件的文件位置
+        :param load_state_dict: 保存的内容是否只是权重。
+        """
         model = self.unwrap_model()
         if isinstance(filepath, Path):
             filepath = str(filepath)
@@ -192,25 +214,15 @@ class PaddleDriver(Driver):
     @rank_zero_call
     def save_checkpoint(self, folder: Path, states: Dict, dataloader, only_state_dict: bool = True, should_save_model: bool = True, **kwargs):
         r"""
-        断点重训的保存函数，该函数会负责保存模型和 optimizers, fp16 的 state_dict；以及模型的保存（若 should_save_model 为 True）
-
-        :param folder: 保存断点重训的状态的文件夹；save 函数应该在下面新增两（一）个文件 的 FASTNLP_CHECKPOINT_FILENAME 文件与
-            FASTNLP_MODEL_FILENAME （如果 should_save_model 为 True ）。把 model 相关的内容放入到 FASTNLP_MODEL_FILENAME 文件中，
-            将传入的 states 以及自身产生其它状态一并保存在 FASTNLP_CHECKPOINT_FILENAME 里面。
-        :param states: 由 trainer 传入的一个字典，其中已经包含了为了实现断点重训所需要保存的其它对象的状态，Driver 应该只需要保存该对象即可，
-            Driver 应该不需要理解该对象，同时在 driver.load_checkpoint() 的时候，需要将 states 返回回去，load() 返回的值与这里的传入的值保持一致。
-        :param dataloader: 正在使用的 dataloader，需要保存里面的状态使得之后可以从当前迭代的位置恢复。
-        :param only_state_dict: 是否只保存模型的参数，当 should_save_model 为 False ，该参数无效。
-        :param should_save_model: 是否应该保存模型，如果为False，Driver 将不负责 model 的保存。
-        :kwargs:
-            * input_spec -- 描述存储模型 ``forward`` 方法的输入；
-            当 ``only_state_dict`` 为 ``False`` 时必须传入，否则加载时会报错。您可以通过 ``InputSpec`` 或者示例 ``Tensor``
-            进行描述。详细的使用方法可以参考 **PaddlePaddle** `关于 paddle.jit.save 函数的文档 <https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/jit/save_cn.html#save>`_；
-        
-        .. todo:
-
-            等 Driver 的文档写完
-
+        断点重训的保存函数，该函数会负责保存 **优化器** 、 **sampler** 和 **fp16** 的状态，以及 **模型** （若 ``should_save_model`` 为 ``True``）
+
+        :param folder: 保存断点重训的状态的文件夹；:meth:`save_checkpoint` 函数应该在该路径下面下面新增名为 ``FASTNLP_CHECKPOINT_FILENAME`` 与
+            ``FASTNLP_MODEL_FILENAME`` （如果 ``should_save_model`` 为 ``True`` ）的文件。把 model 相关的内容放入到 ``FASTNLP_MODEL_FILENAME`` 文件
+            中，将传入的 ``states`` 以及自身产生的其它状态一并保存在 ``FASTNLP_CHECKPOINT_FILENAME`` 里面。
+        :param states: 由 :class:`~fastNLP.core.controllers.Trainer` 传入的一个字典，其中已经包含了为了实现断点重训所需要保存的其它对象的状态。
+        :param dataloader: 正在使用的 dataloader。
+        :param only_state_dict: 是否只保存模型的参数，当 ``should_save_model`` 为 ``False`` ，该参数无效。
+        :param should_save_model: 是否应该保存模型，如果为 ``False`` ，Driver 将不负责 model 的保存。
         """
         # 传入的 dataloader 参数是 trainer 的 dataloader 属性，因为 driver 的所有 dataloader 我们是不会去改变它的，而是通过改变
         #  trainer.dataloader 来改变 dataloader 的状态，从而适配训练或者评测环境；
@@ -272,7 +284,36 @@ class PaddleDriver(Driver):
         logger.debug("Load optimizer state dict.")
 
     def load_checkpoint(self, folder: Path, dataloader, only_state_dict: bool = True, should_load_model: bool = True, **kwargs) -> Dict:
-        
+        r"""
+        断点重训的加载函数，该函数会负责读取数据，并且恢复 **优化器** 、**sampler** 、 **fp16** 的状态和 **模型** （如果 ``should_load_model`` 为 True）以及其它
+        在 :meth:`save_checkpoint` 函数中执行的保存操作，然后将一个 state 字典返回给 :class:`~fastNLP.core.controllers.Trainer` （ 内容为 :meth:`save_checkpoint` 
+        接受到的 ``states`` ）。
+
+        该函数应该在所有 rank 上执行。
+
+        :param folder: 读取该 folder 下的 ``FASTNLP_CHECKPOINT_FILENAME`` 文件与 ``FASTNLP_MODEL_FILENAME``
+            （如果 should_load_model 为True）。
+        :param dataloader: 当前给定 dataloader，需要根据保存的 dataloader 状态合理设置。若该值为 ``None`` ，则不需要返回 ``'dataloader'``
+            以及 ``'batch_idx_in_epoch'`` 这两个值。
+        :param only_state_dict: 是否仅读取模型的 state_dict ，当 ``should_save_model`` 为 ``False`` ，该参数无效。如果为 ``True`` ，说明保存的内容为权重；如果为
+            False 说明保存的是模型，但也是通过当前 Driver 的模型去加载保存的模型的权重，而不是使用保存的模型替换当前模型。
+        :param should_load_model: 是否应该加载模型，如果为 ``False`` ，Driver 将不负责加载模型。若该参数为 ``True`` ，但在保存的状态中没有
+            找到对应的模型状态，则报错。
+        :return: :meth:`save_checkpoint` 函数输入的 ``states`` 内容。除此之外，还返回的内容有：
+
+            * *dataloader* -- 根据传入的 ``dataloader`` 与读取出的状态设置为合理状态的 dataloader。在当前 ``dataloader`` 样本数与读取出的 sampler 样本数
+              不一致时报错。
+            * *batch_idx_in_epoch* -- :class:`int` 类型的数据，表明当前 epoch 进行到了第几个 batch 。请注意，该值不能仅通过保存的数据中读取的，因为前后两次运行的
+              ``batch_size`` 可能有变化，而应该符合以下等式::
+
+                返回的 dataloader 还会产生的 batch 数量 + batch_idx_in_epoch = 原来不断点训练时的 batch 的总数
+              
+              由于 ``返回的 dataloader 还会产生的batch数`` 在 ``batch_size`` 与 ``drop_last`` 参数给定的情况下，无法改变，因此只能通过调整 ``batch_idx_in_epoch``
+              这个值来使等式成立。一个简单的计算原则如下：
+
+                * drop_last 为 ``True`` 时，等同于 floor(sample_in_this_rank/batch_size) - floor(num_left_samples/batch_size)；
+                * drop_last 为 ``False`` 时，等同于 ceil(sample_in_this_rank/batch_size) - ceil(num_left_samples/batch_size)。
+        """
         states = paddle.load(str(folder.joinpath(FASTNLP_CHECKPOINT_FILENAME)))
 
         # 1. 加载 optimizers 的状态；
@@ -333,7 +374,7 @@ class PaddleDriver(Driver):
 
     def get_evaluate_context(self):
         r"""
-        返回一个不计算梯度的环境用来对模型进行评测；
+        返回一个不计算梯度的环境用来对模型进行评测。
 
         :return: 上下文对象 ``paddle.no_grad``;
         """
@@ -342,14 +383,14 @@ class PaddleDriver(Driver):
     @staticmethod
     def move_model_to_device(model: "paddle.nn.Layer", device: Union[str, int, "paddle.CUDAPlace", "paddle.CPUPlace"]):
         r"""
-        用来将模型 ``model`` 转移到指定的设备上；
+        用来将模型 ``model`` 转移到指定的设备上。
 
         .. note::
 
             在 **Paddle** 中使用可能会引起因与设置的设备不一致而产生的问题，请注意。
 
-        :param model: 需要进行转移的模型；
-        :param device: 目标设备；
+        :param model: 需要进行转移的模型。
+        :param device: 目标设备。
         """
         if device is not None:
             model.to(device)
@@ -362,8 +403,8 @@ class PaddleDriver(Driver):
 
             在 **Paddle** 中使用可能会引起因与设置的设备不一致而产生的问题，请注意。
 
-        :param batch: 包含 :class:`paddle.Tensor` 的数据集合，可以是 **List**、**Dict** 等嵌套类型；
-        :return: 移动到指定机器后的 `batch``；
+        :param batch: 包含 :class:`paddle.Tensor` 的数据集合，可以是 **List**、**Dict** 等嵌套类型。
+        :return: 移动到指定机器后的 ``batch``。
         """
         device = _convert_data_device(self.data_device)
         return paddle_move_data_to_device(batch, device)
@@ -387,10 +428,19 @@ class PaddleDriver(Driver):
         random.seed(stdlib_seed)
 
     def set_deterministic_dataloader(self, dataloader):
+        """
+        为了确定性训练要对 ``dataloader`` 进行修改，保证在确定随机数种子后，每次重新训练得到的结果是一样的。 
+        """
         if dataloader.worker_init_fn is None:
             dataloader.worker_init_fn = partial(self.worker_init_function, rank=self.global_rank)
 
     def set_sampler_epoch(self, dataloader: "DataLoader", cur_epoch_idx):
+        r"""
+        对于分布式的 ``sampler``，需要在每一个 ``epoch`` 前设置随机数种子，来保证每一个进程上的 ``shuffle`` 是一样的。
+
+        :param dataloader: 需要设置 ``epoch`` 的 ``dataloader``
+        :param cur_epoch_idx: 当前是第几个 ``epoch``
+        """
         if callable(getattr(dataloader.batch_sampler, "set_epoch", None)):
             dataloader.batch_sampler.set_epoch(cur_epoch_idx)
         elif callable(getattr(dataloader.batch_sampler.sampler, "set_epoch", None)):
@@ -398,7 +448,10 @@ class PaddleDriver(Driver):
 
     @staticmethod
     def get_dataloader_args(dataloader: "DataLoader"):
-
+        """
+        从 ``dataloader`` 中获取参数 ``dataset``, ``batch_sampler``, ``sampler``, ``batch_size``, ``shuffle`` 
+        和 ``drop_last`` 。
+        """
         @dataclass
         class Res:
             dataset: Optional[Dataset] = None

fastNLP/core/drivers/paddle_driver/single_device.py

@@ -40,13 +40,13 @@ class PaddleSingleDriver(PaddleDriver):
     """
     实现了 **PaddlePaddle** 框架下在单卡或 ``cpu`` 环境下训练功能的 **Driver**。
 
-    :param model: 训练时使用的 **PaddlePaddle** 模型；
-    :param device: 训练使用的设备；
-    :param fp16: 是否开启混合精度训练；
+    :param model: 训练时使用的 **PaddlePaddle** 模型
+    :param device: 训练使用的设备
+    :param fp16: 是否开启混合精度训练
     :param paddle_kwargs:
-        * *gradscaler_kwargs* -- 用于 ``fp16=True`` 时，提供给 :class:`paddle.amp.GradScaler` 的参数;
+        * *gradscaler_kwargs* -- 用于 ``fp16=True`` 时，提供给 :class:`paddle.amp.GradScaler` 的参数。
     :kwargs:
-        * wo_auto_param_call (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为；
+        * *model_wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为。
 
         .. note::
 
@@ -155,19 +155,19 @@ class PaddleSingleDriver(PaddleDriver):
 
     def unwrap_model(self):
         """
-        返回训练使用的模型。
+        :return: 训练使用的模型。
         """
         return self.model
 
     @property
     def data_device(self) -> str:
         """
-        :return: 数据和模型所在的设备；
+        :return: 数据和模型所在的设备。
         """
         return self.model_device
 
     def is_distributed(self) -> bool:
         """
-        判断是否为分布式的 **Driver** ，在 ``PaddleSingleDriver`` 中，返回 ``False``。
+        :return 是否为分布式的 **Driver** ，在 ``PaddleSingleDriver`` 中，返回 ``False``。
         """
         return False

fastNLP/core/drivers/paddle_driver/utils.py

@@ -32,6 +32,7 @@ else:
 
 __all__ = [
     "paddle_seed_everything",
+    "optimizer_state_to_device",
 ]
 
 def paddle_seed_everything(seed: int = None, add_global_rank_to_seed: bool = True) -> int:
@@ -40,7 +41,7 @@ def paddle_seed_everything(seed: int = None, add_global_rank_to_seed: bool = Tru
 
     :param seed: 全局随机状态的整数值种子。如果为 ``None`` 则会根据时间戳生成一个种子。
     :param add_global_rank_to_seed: 在分布式训练中，是否在不同 **rank** 中使用不同的随机数。
-        当设置为 ``True`` 时，**FastNLP** 会将种子加上当前的 ``global_rank``。
+        当设置为 ``True`` 时，**fastNLP** 会将种子加上当前的 ``global_rank``。
     """
     max_seed_value = np.iinfo(np.uint32).max
     min_seed_value = np.iinfo(np.uint32).min
@@ -247,6 +248,13 @@ def replace_sampler(dataloader, new_sampler):
     return replace_batch_sampler(dataloader, new_batch_sampler)
 
 def optimizer_state_to_device(state, device):
+    r"""
+    将一个 ``optimizer`` 的 ``state_dict`` 迁移到对应的设备。
+
+    :param state: :func:`optimzier.state_dict` 获取的 state_dictt
+    :param device: 要迁移到的目的设备。
+    :return: 迁移后的新的 state_dict。
+    """
     new_state = {}
     for name, param in state.items():
         if isinstance(param, dict):

fastNLP/core/drivers/torch_driver/__init__.py

@@ -1,17 +1,21 @@
 __all__ = [
-    'TorchDDPDriver',
+    'TorchDriver',
     'TorchSingleDriver',
+    'TorchDDPDriver',
+    'airScaleDriver',
     'DeepSpeedDriver',
-    'TorchDriver',
+    'TorchFSDPDriver',
     'torch_seed_everything',
     'optimizer_state_to_device'
 ]
 
 from .ddp import TorchDDPDriver
 # todo 实现 fairscale 后再将 fairscale 导入到这里；
+from .fairscale import FairScaleDriver
 from .single_device import TorchSingleDriver
 from .torch_driver import TorchDriver
 from .deepspeed import DeepSpeedDriver
+from .torch_fsdp import TorchFSDPDriver
 from .utils import torch_seed_everything, optimizer_state_to_device
 
 

fastNLP/core/drivers/torch_driver/ddp.py

@@ -164,11 +164,11 @@ from .utils import _check_dataloader_args_for_distributed
 
 class TorchDDPDriver(TorchDriver):
     r"""
-    ``TorchDDPDriver`` 通过开启多个进程，让每个进程单独使用一个 gpu 设备来实现分布式训练；
+    ``TorchDDPDriver`` 通过开启多个进程，让每个进程单独使用一个 gpu 设备来实现分布式训练。
 
     .. note::
 
-        您在绝大多数情况下不需要自己使用到该类，通过向 ``Trainer`` 传入正确的参数，您可以方便快速地部署您的分布式训练；
+        您在绝大多数情况下不需要自己使用到该类，通过向 ``Trainer`` 传入正确的参数，您可以方便快速地部署您的分布式训练。
 
         ``TorchDDPDriver`` 目前支持的三种启动方式：
 
@@ -229,18 +229,24 @@ class TorchDDPDriver(TorchDriver):
             通过运行 ``python -m torch.distributed.launch --nproc_per_node 2 train.py`` 启动；
 
         注意多机的启动强制要求用户在每一台机器上使用 ``python -m torch.distributed.launch`` 启动；因此我们不会在 ``TorchDDPDriver`` 中保存
-        任何当前有多少台机器的信息；
+        任何当前有多少台机器的信息。
 
-    :param model: 传入给 ``Trainer`` 的 ``model`` 参数；
-    :param parallel_device: 用于分布式训练的 ``gpu`` 设备；
-    :param is_pull_by_torch_run: 标志当前的脚本的启动是否由 ``python -m torch.distributed.launch`` 启动的；
-    :param fp16: 是否开启 fp16 训练；
+    :param model: 传入给 ``Trainer`` 的 ``model`` 参数
+    :param parallel_device: 用于分布式训练的 ``gpu`` 设备
+    :param is_pull_by_torch_run: 标志当前的脚本的启动是否由 ``python -m torch.distributed.launch`` 启动的
+    :param fp16: 是否开启 fp16 训练
     :param torch_kwargs: 
         * *ddp_kwargs* -- 用于在使用 ``TorchDDPDriver`` 时指定 ``DistributedDataParallel`` 初始化时的参数；例如传入
-            {'find_unused_parameters': True} 来解决有参数不参与前向运算导致的报错等；
-        * *set_grad_to_none* -- 是否在训练过程中在每一次 optimizer 更新后将 grad 置为 None；
-        * *non_blocking* -- 表示用于 pytorch 的 tensor 的 to 方法的参数 non_blocking；
-        * *gradscaler_kwargs* -- 用于 fp16=True 时，提供给 ``torch.amp.cuda.GradScaler`` 的参数;
+          ``{'find_unused_parameters': True}`` 来解决有参数不参与前向运算导致的报错等
+        * *set_grad_to_none* -- 是否在训练过程中在每一次 optimizer 更新后将 grad 置为 ``None``
+        * *non_blocking* -- 表示用于 :meth:`torch.Tensor.to` 方法的参数 non_blocking
+        * *gradscaler_kwargs* -- 用于 ``fp16=True`` 时，提供给 :class:`torch.amp.cuda.GradScaler` 的参数
+    :kwargs:
+        * *wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为
+
+        .. note::
+
+            关于该参数的详细说明，请参见 :class:`~fastNLP.core.controllers.Trainer` 中的描述；函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`。
     """
 
     def __init__(
@@ -309,9 +315,9 @@ class TorchDDPDriver(TorchDriver):
         self.world_size = None  # int(os.environ.get("WORLD_SIZE"))  len(self.parallel_device)
         self.global_rank = 0
 
-        self._ddp_kwargs = self._torch_kwargs.get("ddp_kwargs", {})
-        check_user_specific_params(self._ddp_kwargs, DistributedDataParallel.__init__, DistributedDataParallel.__name__)
-        if len(self.model._buffers) != 0 and self._ddp_kwargs.get("broadcast_buffers", None) is None:
+        self._fsdp_kwargs = self._torch_kwargs.get("ddp_kwargs", {})
+        check_user_specific_params(self._fsdp_kwargs, DistributedDataParallel.__init__, DistributedDataParallel.__name__)
+        if len(self.model._buffers) != 0 and self._fsdp_kwargs.get("broadcast_buffers", None) is None:
             logger.info("Notice your model has buffers and you are using `TorchDDPDriver`, but you do not set "
                         "'broadcast_buffers' in your trainer. Cause in most situations, this parameter can be set"
                         " to 'False' to avoid redundant data communication between different processes.")
@@ -329,7 +335,7 @@ class TorchDDPDriver(TorchDriver):
         r"""
         准备分布式环境，该函数主要做以下两件事情：
 
-            1. 开启多进程，每个 gpu 设备对应单独的一个进程；
+            1. 开启多进程，每个 ``gpu`` 设备对应单独的一个进程；
             2. 每个进程将模型迁移到自己对应的 ``gpu`` 设备上；然后使用 ``DistributedDataParallel`` 包裹模型；
         """
         if self._has_setup:
@@ -381,8 +387,6 @@ class TorchDDPDriver(TorchDriver):
                 self.global_rank = dist.get_rank()
 
         if not self.outside_ddp:
-            torch.cuda.set_device(self.model_device)
-            self.model.to(self.model_device)
             self.configure_ddp()
 
         self.barrier()
@@ -400,11 +404,13 @@ class TorchDDPDriver(TorchDriver):
         self._pids = self.tensor_to_numeric(self._pids)
 
     def configure_ddp(self):
+        torch.cuda.set_device(self.model_device)
+        self.model.to(self.model_device)
         if not isinstance(self.model, DistributedDataParallel):
             self.model = DistributedDataParallel(
                 # 注意这里的 self.model_device 是 `torch.device` type，因此 self.model_device.index；
                 _DDPWrappingModel(self.model), device_ids=[self.model_device.index],
-                **self._ddp_kwargs
+                **self._fsdp_kwargs
             )
             self._has_ddpwrapped = True
 
@@ -450,10 +456,16 @@ class TorchDDPDriver(TorchDriver):
 
     @property
     def master_address(self) -> str:
+        """
+        分布式训练中的地址 ``MASTER_ADDR``
+        """
         return os.environ.get("MASTER_ADDR", "127.0.0.1")
 
     @property
     def master_port(self) -> str:
+        """
+        分布式训练使用的端口 ``MASTER_PORT``
+        """
         if self.outside_ddp:
             return os.environ.get("MASTER_PORT")
         if self._master_port is None:
@@ -462,6 +474,9 @@ class TorchDDPDriver(TorchDriver):
 
     @property
     def world_size(self) -> int:
+        """
+        分布式训练的进程总数 ``WORLD_SIZE``
+        """
         return self._world_size
 
     @world_size.setter
@@ -470,6 +485,9 @@ class TorchDDPDriver(TorchDriver):
 
     @property
     def global_rank(self) -> int:
+        """
+        当前进程的全局编号 ``global_rank``
+        """
         return self._global_rank
 
     @global_rank.setter
@@ -478,6 +496,9 @@ class TorchDDPDriver(TorchDriver):
 
     @property
     def local_rank(self) -> int:  # 这个不会受到 all_rank_call_context 的影响
+        """
+        当前进程的局部编号 ``local_rank``
+        """
         return int(os.environ.get("LOCAL_RANK", 0))
 
     @property
@@ -505,6 +526,12 @@ class TorchDDPDriver(TorchDriver):
                     raise RuntimeError(f"The `{fn}` attribute of model is not `Callable`.")
                 return fn, None
             elif fn in {"train_step", "evaluate_step"}:
+
+                logger.warning("\n\nfucking hei\n\n")
+                print(model)
+                print("\n\n")
+                print(type(model))
+                print("\n\n")
                 return model, model.forward
             else:
                 raise RuntimeError(f"There is no `{fn}` method in your model.")
@@ -609,20 +636,20 @@ class TorchDDPDriver(TorchDriver):
 
     def is_global_zero(self):
         r"""
-        :return: 返回当前的进程是否在全局上是进程 0 ；
+        :return: 当前的进程是否在全局上是进程 0 。
         """
         return self.global_rank == 0
 
     def get_model_no_sync_context(self):
         r"""
-        :return: 返回一个 ``context`` 上下文环境，用于关闭各个进程之间的同步；
+        :return: 一个 ``context`` 上下文环境，用于关闭各个进程之间的同步。
         """
         # 注意此时的 model 是 "DistributedDataParallel" 对象；
         return self.model.no_sync
 
     def unwrap_model(self):
         r"""
-        :return: 返回没有经过 ``DistributedDataParallel`` 包裹的原始模型；
+        :return: 没有经过 ``DistributedDataParallel`` 包裹的原始模型。
         """
         _module = self.model.module
         if isinstance(_module, _DDPWrappingModel):
@@ -632,34 +659,33 @@ class TorchDDPDriver(TorchDriver):
 
     def get_local_rank(self) -> int:
         r"""
-        :return: 返回当前进程局部的进程编号；
+        :return: 当前进程局部的进程编号。
         """
         return self.local_rank
 
     def barrier(self):
         r"""
-        通过使用该函数来使得各个进程之间同步操作；
+        通过使用该函数来使得各个进程之间同步操作。
         """
         if int(os.environ.get(FASTNLP_NO_SYNC, 0)) < 1:  # 当 FASTNLP_NO_SYNC 小于 1 时实际执行
             torch.distributed.barrier(async_op=False)
 
     def is_distributed(self):
         r"""
-        :return: 返回当前使用的 driver 是否是分布式的 driver，对于 ``TorchDDPDriver`` 来说，该函数一定返回 ``True``；
+        :return: 当前使用的 driver 是否是分布式的 driver，对于 ``TorchDDPDriver`` 来说，该函数一定返回 ``True``。
         """
         return True
 
     def broadcast_object(self, obj, src: int = 0, group=None, **kwargs):
         r"""
-        从 src 端将 obj 对象（可能是 tensor ，可能是 object ）发送到 dst 处。如果是非 tensor 的对象会尝试使用 pickle 进行打包进行
-        传输，然后再 dst 处再加载回来。仅在分布式的 driver 中有实际意义。
+        从 ``src`` 端将 ``obj`` 对象（可能是 tensor ，可能是 object ）广播到其它进程。如果是非 tensor 的对象会尝试使用 pickle 进行打包进行
+        传输，然后在接收处处再加载回来。仅在分布式的 driver 中有实际意义。
 
         :param obj: obj，可能是 Tensor 或 嵌套类型的数据
-        :param int src: source 的 global rank 。
-        :param int dst: target 的 global rank，可以是多个目标 rank
-        :param group: 所属的 group
-        :return: 如果当前不是分布式 driver 直接返回输入的 obj 。如果当前 rank 是接收端（其 global rank 包含在了 dst 中），则返回
-            接收到的参数；如果是 source 端则返回发射的内容；既不是发送端、又不是接收端，则返回 None 。
+        :param src: 发送方的 ``global_rank``
+        :param group: 进程所在的通信组
+        :return: 如果当前 rank 是接收端，则返回接收到的参数；如果是 source 端则返回发送的内容。如果环境变量 ``FASTNLP_NO_SYNC`` 为 **2** 则
+            返回 ``None``
         """
         if int(os.environ.get(FASTNLP_NO_SYNC, 0)) == 2:  # 如果 FASTNLP_NO_SYNC == 2 直接返回。
             return
@@ -667,27 +693,24 @@ class TorchDDPDriver(TorchDriver):
 
     def all_gather(self, obj, group) -> List:
         r"""
-        将 obj 互相传送到其它所有的 rank 上，其中 obj 可能是 Tensor，也可能是嵌套结构的 object 。如果不是基础类型的数据，尝试通过
+        将 ``obj`` 互相传送到其它所有的 rank 上，其中 ``obj`` 可能是 Tensor，也可能是嵌套结构的 object 。如果不是基础类型的数据，将会尝试通过
         pickle 进行序列化，接收到之后再反序列化。
 
         example::
 
-            obj = {
-                'a': [1, 1],
-                'b': [[1, 2], [1, 2]],
-                'c': {
-                    'd': [1, 2]
-                }
-            }
-            ->
-            [
-                {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
-                {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
-            ]
-
-        :param obj: 需要传输的对象，在每个rank上都应该保持相同的结构。
-        :param group:
-        :return:
+            >>> # rank 0
+            >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 1}}
+            >>> # rank 1
+            >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+            >>> # after all_gather():
+            >>> result = [
+                    {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
+                    {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+                ]
+
+        :param obj: 需要传输的对象，在每个 rank 上都应该保持相同的结构。
+        :param group: 进程所在的通信组。
+        :return: 所有 rank 发送的 ``obj`` 聚合在一起的内容；如果环境变量 ``FASTNLP_NO_SYNC`` 为 **2** 则不会执行，直接返回 ``[obj]`` 。
         """
         if int(os.environ.get(FASTNLP_NO_SYNC, 0)) == 2:  # 如果 FASTNLP_NO_SYNC 表示不执行
             return [obj]
@@ -701,7 +724,7 @@ class TorchDDPDriver(TorchDriver):
 def find_free_network_port() -> str:
     """
     在 localhost 上找到一个空闲端口；
-    当我们不想连接到真正的主节点但必须设置“MASTER_PORT”环境变量时在单节点训练中很有用；
+    当我们不想连接到真正的主节点但必须设置“MASTER_PORT”环境变量时在单节点训练中很有用。
     """
     s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
     s.bind(("", 0))

fastNLP/core/drivers/torch_driver/deepspeed.py

@@ -83,26 +83,36 @@ class DeepSpeedDriver(TorchDDPDriver):
                 )
                 trainer.run()
 
-            通过运行 ``deepspeed train.py`` 启动；
+            通过运行 ``deepspeed train.py`` 启动。
 
-    :param model: 传入给 ``Trainer`` 的 ``model`` 参数；
-    :param parallel_device: 用于分布式训练的 ``gpu`` 设备；
-    :param is_pull_by_torch_run: 标志当前的脚本的启动是否由 ``python -m torch.distributed.launch`` 启动的；
-    :param fp16: 是否开启 fp16 训练；
+    :param model: 传入给 ``Trainer`` 的 ``model`` 参数。
+    :param parallel_device: 用于分布式训练的 ``gpu`` 设备。
+    :param is_pull_by_torch_run: 标志当前的脚本的启动是否由 ``python -m torch.distributed.launch`` 启动的。
+    :param fp16: 是否开启 fp16 训练。
     :param deepspeed_kwargs: 
         * *strategy* -- 使用 ZeRO 优化的策略，默认为 ``deepspeed``；目前仅支持以下值：
 
             * ``deepspeed`` -- 使用 ZeRO 的第二阶段，等同于 ``deepspeed_stage_2``；
             * ``deepspeed_stage_1`` -- 使用 ZeRO 的第一阶段，仅将 ``optimizer`` 的状态分散到不同设备上；
-            * ``deepspeed_stage_2`` -- 使用 ZeRO 的第二阶段，将 ``optimizer`` 和**梯度**分散到不同设备上；
+            * ``deepspeed_stage_2`` -- 使用 ZeRO 的第二阶段，将 ``optimizer`` 和 **梯度** 分散到不同设备上；
             * ``deepspeed_stage_2_offload`` -- 使用 ZeRO 的第二阶段，并且借助 cpu 的内存来进一步节约显存；
-            * ``deepspeed_stage_3`` -- 使用 ZeRO 的第三阶段，将 ``optimizer`` 、**梯度**和**模型**分散到不同设备上；
+            * ``deepspeed_stage_3`` -- 使用 ZeRO 的第三阶段，将 ``optimizer`` 、**梯度** 和 **模型** 分散到不同设备上；
             * ``deepspeed_stage_3_offload`` -- 使用 ZeRO 的第三阶段，并且借助 cpu 的内存来进一步节约显存；
             * ``deepspeed_stage_3_offload_nvme`` -- 使用 ZeRO 的第三阶段，并且借助 NVMe 硬盘来进一步节约显存；
-        * *logging_level* -- ``deepspeed`` 库的日志等级，默认为 **logging.ERROR**；
+        * *logging_level* -- ``deepspeed`` 库的日志等级，默认为 **logging.ERROR**。
         * *config* -- ``deepspeed`` 的各项设置；**FastNLP** 允许用户传入自己的设置以增强灵活性，但这会使参数
           中的 ``optimizer`` 、``strategy`` 、 ``fp16`` 等失效，即当这个参数存在时，**FastNLP** 会用该参数覆盖
-          其它的设置；
+          其它的设置。
+    :kwargs:
+        * *accumulation_steps* -- 即在 :class:`~fastNLP.core.controllers.Trainer` 传入的 ``accumulation_steps`` 。 deepspeed 会将 ``config`` 的
+          ``gradient_accumulation_steps`` 设置为该值。
+        * *train_dataloader* -- 即在 :class:`~fastNLP.core.controllers.Trainer` 传入的 ``train_dataloader`` 。 ``deepspeed`` 需要通过它来获取
+          数据的 ``batch_size`` 用于设置 ``train_micro_batch_size_per_gpu`` 。如果没有传入的话，则会设置为 **1** 。
+        * *wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为
+
+        .. note::
+
+            关于该参数的详细说明，请参见 :class:`~fastNLP.core.controllers.Trainer` 中的描述；函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`。
     """
     # TODO fp16 load_config
     def __init__(
@@ -357,28 +367,38 @@ class DeepSpeedDriver(TorchDDPDriver):
                 self.config["amp"] = {"enabled": True, "opt_level": "O1"}
 
     def zero_grad(self):
+        """
+        进行梯度置零操作；由于 :meth:`DeepSpeedEngine.step` 包含了 :meth:`zero_step` 的功能，因此该接口实际无意义。
+        """
         # DeepSpeedEngine.step 包含了 zero_grad 功能
         pass
 
     def backward(self, loss):
+        """
+        对 ``loss`` 进行反向传播
+        """
         self.model.backward(loss)
 
     def step(self):
+        """
+        更新模型的参数
+        """
         self.model.step()
 
     def get_model_no_sync_context(self):
         r"""
-        :return: 返回一个 ``context`` 上下文环境，用于关闭各个进程之间的同步；在 ``deepspeed`` 中，返回一个空的上下文
+        :return: 一个 ``context`` 上下文环境，用于关闭各个进程之间的同步；在 ``deepspeed`` 中，返回一个空的上下文
         """
         # 注意此时的 model 是 "DistributedDataParallel" 对象；
         return nullcontext
 
     def save_model(self, filepath: Union[str, Path], only_state_dict: bool = False, **kwargs):
         """
-        保存当前 driver 的模型到 folder 下。
+        保存的模型到 ``filepath`` 中。
 
-        :param filepath: 保存到哪个文件夹；
-        :param only_state_dict: 是否只保存权重；在 ``DeepSpeedDriver`` 中该参数无效；
+        :param filepath: 文件路径
+        :param only_state_dict: 是否只保存权重；在 ``DeepSpeedDriver`` 中该参数无效。
+        :param kwargs: 需要传入 **deepspeed** 模型 :meth:`save_checkpoint` 的其它参数。
         :return:
         """
         # deepspeed engine 要求在每个 rank 都调用 save_checkpoint，故去掉了 rank_zero_call 装饰器
@@ -398,11 +418,11 @@ class DeepSpeedDriver(TorchDDPDriver):
 
     def load_model(self, filepath: Union[Path, str], only_state_dict: bool = False, **kwargs):
         """
-        从 folder 中加载权重并赋值到当前 driver 的模型上。
+        从 ``filepath`` 中加载权重并赋值到当前 driver 的模型上。
 
         :param filepath: 加载权重或模型的路径
-        :param load_state_dict: 保存的内容是否只是权重；在 ``DeepSpeedDriver`` 中该参数无效；
-        :param kwargs:
+        :param load_state_dict: 保存的内容是否只是权重；在 ``DeepSpeedDriver`` 中该参数无效。
+        :param kwargs: 需要传入 **deepspeed** 模型 :meth:`load_checkpoint` 的其它参数。
         :return:
         """
         if not only_state_dict:
@@ -411,6 +431,17 @@ class DeepSpeedDriver(TorchDDPDriver):
         self.model.load_checkpoint(filepath, **kwargs)
 
     def save_checkpoint(self, folder: Path, states: Dict, dataloader, only_state_dict: bool = True, should_save_model: bool = True, **kwargs):
+        r"""
+        断点重训的保存函数，该函数会负责保存 **优化器** 、 **sampler** 和 **fp16** 的状态，以及 **模型** （若 ``should_save_model`` 为 ``True``）
+
+        :param folder: 保存断点重训的状态的文件夹；:meth:`save_checkpoint` 函数应该在该路径下面下面新增名为 ``FASTNLP_CHECKPOINT_FILENAME`` 与
+            ``FASTNLP_MODEL_FILENAME`` （如果 ``should_save_model`` 为 ``True`` ）的文件。把 model 相关的内容放入到 ``FASTNLP_MODEL_FILENAME`` 文件
+            中，将传入的 ``states`` 以及自身产生其它状态一并保存在 ``FASTNLP_CHECKPOINT_FILENAME`` 里面。
+        :param states: 由 :class:`~fastNLP.core.controllers.Trainer` 传入的一个字典，其中已经包含了为了实现断点重训所需要保存的其它对象的状态。
+        :param dataloader: 正在使用的 dataloader。
+        :param only_state_dict: 是否只保存模型的参数，当 ``should_save_model`` 为 ``False`` ，该参数无效。
+        :param should_save_model: 是否应该保存模型，如果为 ``False`` ，Driver 将不负责 model 的保存。
+        """
         # deepspeed engine 要求在每个 rank 都调用 save_checkpoint，故去掉了 rank_zero_call 装饰器
         # 1. 保存 sampler 的状态
         num_consumed_batches = states.pop('num_consumed_batches')
@@ -425,6 +456,36 @@ class DeepSpeedDriver(TorchDDPDriver):
                                     client_state=states)
 
     def load_checkpoint(self, folder: Path, dataloader, only_state_dict: bool = True, should_load_model: bool = True, **kwargs) -> Dict:
+        r"""
+        断点重训的加载函数，该函数会负责读取数据，并且恢复 **优化器** 、**sampler** 、 **fp16** 的状态和 **模型** （如果 ``should_load_model`` 为 True）以及其它
+        在 :meth:`save_checkpoint` 函数中执行的保存操作，然后将一个 state 字典返回给 :class:`~fastNLP.core.controllers.Trainer` （ 内容为 :meth:`save_checkpoint` 
+        接受到的 ``states`` ）。
+
+        该函数应该在所有 rank 上执行。
+
+        :param folder: 读取该 folder 下的 ``FASTNLP_CHECKPOINT_FILENAME`` 文件与 ``FASTNLP_MODEL_FILENAME``
+            （如果 should_load_model 为True）。
+        :param dataloader: 当前给定 dataloader，需要根据保存的 dataloader 状态合理设置。若该值为 ``None`` ，则不需要返回 ``'dataloader'``
+            以及 ``'batch_idx_in_epoch'`` 这两个值。
+        :param only_state_dict: 是否仅读取模型的 state_dict ，当 ``should_save_model`` 为 ``False`` ，该参数无效。如果为 ``True`` ，说明保存的内容为权重；如果为
+            False 说明保存的是模型，但也是通过当前 Driver 的模型去加载保存的模型的权重，而不是使用保存的模型替换当前模型。
+        :param should_load_model: 是否应该加载模型，如果为 ``False`` ，Driver 将不负责加载模型。若该参数为 ``True`` ，但在保存的状态中没有
+            找到对应的模型状态，则报错。
+        :return: :meth:`save_checkpoint` 函数输入的 ``states`` 内容。除此之外，还返回的内容有：
+
+            * *dataloader* -- 根据传入的 ``dataloader`` 与读取出的状态设置为合理状态的 dataloader。在当前 ``dataloader`` 样本数与读取出的 sampler 样本数
+              不一致时报错。
+            * *batch_idx_in_epoch* -- :class:`int` 类型的数据，表明当前 epoch 进行到了第几个 batch 。请注意，该值不能仅通过保存的数据中读取的，因为前后两次运行的
+              ``batch_size`` 可能有变化，而应该符合以下等式::
+
+                返回的 dataloader 还会产生的 batch 数量 + batch_idx_in_epoch = 原来不断点训练时的 batch 的总数
+              
+              由于 ``返回的 dataloader 还会产生的batch数`` 在 ``batch_size`` 与 ``drop_last`` 参数给定的情况下，无法改变，因此只能通过调整 ``batch_idx_in_epoch``
+              这个值来使等式成立。一个简单的计算原则如下：
+
+                * drop_last 为 ``True`` 时，等同于 floor(sample_in_this_rank/batch_size) - floor(num_left_samples/batch_size)；
+                * drop_last 为 ``False`` 时，等同于 ceil(sample_in_this_rank/batch_size) - ceil(num_left_samples/batch_size)。
+        """
         # 1. 加载模型状态；
         if not should_load_model:
             logger.rank_zero_warning("Loading checkpoint without model is not allowed for `DeepSpeedDriver`, "
@@ -442,4 +503,7 @@ class DeepSpeedDriver(TorchDDPDriver):
 
     @property
     def stage_3(self) -> bool:
+        """
+        判断是否为第三阶段的 ZeRO 优化
+        """
         return self.config.get("zero_optimization") and self.config.get("zero_optimization").get("stage") == 3

fastNLP/core/drivers/torch_driver/dist_utils.py

@@ -22,6 +22,7 @@ if _NEED_IMPORT_TORCH:
 
 from fastNLP.core.utils import apply_to_collection
 
+__all__ = []
 
 def _validate_output_list_for_rank(my_rank, dst, gather_list):
     if dst == my_rank:
@@ -148,7 +149,7 @@ def send_recv_object(obj, src, cur_rank, device, group=None, tag=0):
     r"""
     pytorch 中的单点对多点的分发函数；
 
-    例如将进程 0 上的对象 object 分发到其它进程上；
+    例如将进程 0 上的对象 object 分发到其它进程上。
 
     Example::
 
@@ -158,11 +159,11 @@ def send_recv_object(obj, src, cur_rank, device, group=None, tag=0):
 
         send_recv_object(object, 0, cur_rank, local_device)
 
-    :param obj: 一个可以序列化的 python 对象；
-    :param src: 从哪一个 rank 上发送到其它 rank；
-    :param cur_rank: 当前的进程的 rank 序号；
-    :param device: 当前的进程所在的设备；
-    :param group: 通信组，默认为 None；
+    :param obj: 一个可以序列化的 python 对象。
+    :param src: 从哪一个 rank 上发送到其它 rank。
+    :param cur_rank: 当前的进程的 rank 序号。
+    :param device: 当前的进程所在的设备。
+    :param group: 通信组，默认为 None。
     :param tag: 将发送与远程接收匹配的标记；
     :return:
     """
@@ -198,18 +199,15 @@ def fastnlp_torch_all_gather(obj: Any, device=None, group=DEFAULT_TORCH_GROUP) -
 
     example::
 
-        obj = {
-            'a': [1, 1],
-            'b': [[1, 2], [1, 2]],
-            'c': {
-                'd': [1, 2]
-            }
-        }
-        ->
-        [
-            {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
-            {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
-        ]
+        >>> # rank 0
+        >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 1}}
+        >>> # rank 1
+        >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+        >>> # after all_gather():
+        >>> result = [
+                {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
+                {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+            ]
 
     :param obj: 任意结构的数据，如果为 tensor ，需要保证每个显卡上的 tensor 的形状是一样的。如果传入的是非 tensor 对象都将直接进行
         序列化之后进行传输。

fastNLP/core/drivers/torch_driver/fairscale.py

@@ -29,6 +29,29 @@ from .utils import optimizer_state_to_device
 
 
 class FairScaleDriver(TorchDDPDriver):
+    """
+    实现 ``fairscale`` 功能的 ``Driver`` 。
+
+    :param model: 传入给 ``Trainer`` 的 ``model`` 参数。
+    :param parallel_device: 用于分布式训练的 ``gpu`` 设备。
+    :param is_pull_by_torch_run: 标志当前的脚本的启动是否由 ``python -m torch.distributed.launch`` 启动的。
+    :param fp16: 是否开启 fp16 训练。
+    :param fairscale_kwargs:
+
+        * *oss_kwargs* --
+        * *sdp_kwargs* --
+        * *fsdp_kwargs* --
+        * *ddp_kwargs* --
+        * *set_grad_to_none* -- 是否在训练过程中在每一次 optimizer 更新后将 grad 置为 ``None``
+        * *non_blocking* -- 表示用于 :meth:`torch.Tensor.to` 方法的参数 non_blocking
+        * *gradscaler_kwargs* -- 用于 ``fp16=True`` 时，提供给 :class:`torch.amp.cuda.GradScaler` 的参数
+    :kwargs:
+        * *wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为
+
+        .. note::
+
+            关于该参数的详细说明，请参见 :class:`~fastNLP.core.controllers.Trainer` 中的描述；函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`。
+    """
     def __init__(
             self,
             model,
@@ -301,6 +324,6 @@ class FairScaleDriver(TorchDDPDriver):
 
     def unwrap_model(self):
         r"""
-        :return: 返回原本的模型，例如没有被 ``DataParallel`` 包裹；
+        :return: 原本的模型，例如没有被 ``DataParallel`` 包裹；
         """
         return self.model.module.model

fastNLP/core/drivers/torch_driver/initialize_torch_driver.py

@@ -9,6 +9,7 @@ from .single_device import TorchSingleDriver
 from .ddp import TorchDDPDriver
 from .fairscale import FairScaleDriver
 from .deepspeed import DeepSpeedDriver
+from .torch_fsdp import TorchFSDPDriver
 from fastNLP.core.log import logger
 from fastNLP.envs import FASTNLP_BACKEND_LAUNCH
 from pkg_resources import parse_version
@@ -21,11 +22,16 @@ def initialize_torch_driver(driver: str, device: Optional[Union[str, "torch.devi
     r"""
     用来根据参数 ``driver` 和 ``device`` 来确定并且初始化一个具体的 ``Driver`` 实例然后返回回去；
 
-    :param driver: 该参数的值应为以下之一：``["torch", "fairscale", "deepspeed"]``；
-    :param device: 该参数的格式与 ``Trainer`` 对参数 ``device`` 的要求一致；
-    :param model: 训练或者评测的具体的模型；
+    :param driver: 该参数的值应为以下之一：``["torch", "fairscale", "deepspeed", "torch_fsdp"]``
+    :param device: 该参数的格式与 ``Trainer`` 对参数 ``device`` 的要求一致
+    :param model: 训练或者评测的具体的模型
 
-    :return: 返回一个 :class:`~fastNLP.core.TorchSingleDriver` 或 :class:`~fastNLP.core.TorchDDPDriver` 实例；
+    :return: 下列类型之一的实例：
+        * :class:`~fastNLP.core.drivers.torch_driver.TorchSingleDriver` 
+        * :class:`~fastNLP.core.drivers.torch_driver.TorchDDPDriver` 
+        * :class:`~fastNLP.core.drivers.torch_driver.DeepSpeedDriver`
+        * :class:`~fastNLP.core.drivers.torch_driver.FairScaleDriver`
+        * :class:`~fastNLP.core.drivers.torch_driver.TorchFSDPDriver`
     """
     if parse_version(torch.__version__) < parse_version('1.6'):
         raise RuntimeError(f"Pytorch(current version:{torch.__version__}) need to be older than 1.6.")
@@ -45,7 +51,7 @@ def initialize_torch_driver(driver: str, device: Optional[Union[str, "torch.devi
             return TorchDDPDriver(model, torch.device(f"cuda:{os.environ['LOCAL_RANK']}"),
                                   is_pull_by_torch_run=True, **kwargs)
 
-    if driver not in {"torch", "fairscale", "deepspeed"}:
+    if driver not in {"torch", "fairscale", "deepspeed", "torch_fsdp"}:
         raise ValueError("Parameter `driver` can only be one of these values: ['torch', 'fairscale'].")
 
     _could_use_device_num = torch.cuda.device_count()
@@ -95,4 +101,12 @@ def initialize_torch_driver(driver: str, device: Optional[Union[str, "torch.devi
             logger.warning_once("Notice you are using `deepspeed`, but the `device` is only one gpu.")
             return DeepSpeedDriver(model, [device], **kwargs)
         else:
-            return DeepSpeedDriver(model, device, **kwargs)
+            return DeepSpeedDriver(model, device, **kwargs)
+    elif driver == "torch_fsdp":
+        if not isinstance(device, List):
+            if device.type == 'cpu':
+                raise ValueError("You are using `torch_fsdp` driver, but your chosen `device` is 'cpu'.")
+            logger.warning_once("Notice you are using `torch_fsdp`, but the `device` is only one gpu.")
+            return TorchFSDPDriver(model, [device], **kwargs)
+        else:
+            return TorchFSDPDriver(model, device, **kwargs)

fastNLP/core/drivers/torch_driver/single_device.py

@@ -26,19 +26,25 @@ from fastNLP.core.log import logger
 
 class TorchSingleDriver(TorchDriver):
     r"""
-    ``TorchSingleDriver`` 是用于 cpu 和 单卡 gpu 运算的 ``driver``；
+    ``TorchSingleDriver`` 是用于 cpu 和 单卡 gpu 运算的 ``driver``。
 
     .. note::
 
-        如果您希望使用 ``DataParallel`` 来训练您的模型，您应当自己在 ``Trainer`` 初始化之前初始化好 ``DataParallel``，然后将其传入 ``Trainer`` 中；
+        如果您希望使用 ``DataParallel`` 来训练您的模型，您应当自己在 ``Trainer`` 初始化之前初始化好 ``DataParallel``，然后将其传入 ``Trainer`` 中。
 
-    :param model: 传入给 ``Trainer`` 的 ``model`` 参数；
-    :param device: torch.device，当前进程所使用的设备；
-    :param fp16: 是否开启 fp16；
+    :param model: 传入给 ``Trainer`` 的 ``model`` 参数
+    :param device: torch.device，当前进程所使用的设备
+    :param fp16: 是否开启 fp16
     :param torch_kwargs:
-        * *set_grad_to_none* -- 是否在训练过程中在每一次 optimizer 更新后将 grad 置为 None；
-        * *non_blocking* -- 表示用于 pytorch 的 tensor 的 to 方法的参数 non_blocking；
-        * *gradscaler_kwargs* -- 用于 fp16=True 时，提供给 ``torch.amp.cuda.GradScaler`` 的参数;
+        * *set_grad_to_none* -- 是否在训练过程中在每一次 optimizer 更新后将 grad 置为 ``None``
+        * *non_blocking* -- 表示用于 :meth:`torch.Tensor.to` 方法的参数 non_blocking
+        * *gradscaler_kwargs* -- 用于 ``fp16=True`` 时，提供给 :class:`torch.amp.cuda.GradScaler` 的参数
+    :kwargs:
+        * *wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为
+
+        .. note::
+
+            关于该参数的详细说明，请参见 :class:`~fastNLP.core.controllers.Trainer` 中的描述；函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`。
     """
 
     def __init__(self, model, device: "torch.device", fp16: bool = False, torch_kwargs: Dict = None, **kwargs):
@@ -69,7 +75,7 @@ class TorchSingleDriver(TorchDriver):
 
     def setup(self):
         r"""
-        将模型迁移到相应的设备上；
+        将模型迁移到相应的设备上。
         """
         if self.model_device is not None:
             self.model.to(self.model_device)
@@ -153,7 +159,7 @@ class TorchSingleDriver(TorchDriver):
 
     def unwrap_model(self):
         r"""
-        :return: 返回原本的模型，例如没有被 ``DataParallel`` 包裹；
+        :return: 原本的模型，该函数可以取出被 ``DataParallel`` 包裹的模型
         """
         if isinstance(self.model, torch.nn.DataParallel) or \
                 isinstance(self.model, torch.nn.parallel.DistributedDataParallel):
@@ -164,12 +170,12 @@ class TorchSingleDriver(TorchDriver):
     @property
     def data_device(self):
         r"""
-        注意单卡模式下使用 ``driver.data_device`` 等价于使用 ``driver.model_device``；
+        数据和模型所在的设备
         """
         return self.model_device
 
     def is_distributed(self):
         r"""
-        :return: 返回当前使用的 driver 是否是分布式的 driver，对于 ``TorchSingleDriver`` 来说直接返回 ``False``；
+        :return: 当前使用的 driver 是否是分布式的 driver，对于 ``TorchSingleDriver`` 来说直接返回 ``False``
         """
         return False

fastNLP/core/drivers/torch_driver/torch_driver.py

@@ -27,7 +27,7 @@ from .utils import optimizer_state_to_device
 from fastNLP.core.drivers.driver import Driver
 from fastNLP.core.drivers.torch_driver.utils import _build_fp16_env, DummyGradScaler
 from fastNLP.core.utils import apply_to_collection, torch_move_data_to_device
-from fastNLP.envs import  rank_zero_call
+from fastNLP.envs import rank_zero_call
 from fastNLP.envs import FASTNLP_GLOBAL_RANK, FASTNLP_MODEL_FILENAME, FASTNLP_CHECKPOINT_FILENAME
 from fastNLP.core.log import logger
 from fastNLP.core.samplers import ReproducibleBatchSampler, ReproducibleSampler, ReproduceBatchSampler, RandomSampler
@@ -36,18 +36,18 @@ from fastNLP.core.dataloaders import OverfitDataLoader
 
 class TorchDriver(Driver):
     r"""
-    专属于 ``pytorch`` 的 ``driver``，是 ``TorchSingleDriver`` 和 ``TorchDDPDriver`` 的父类；
+    专属于 ``pytorch`` 的 ``driver``，是 ``TorchSingleDriver`` 和 ``TorchDDPDriver`` 的父类。
 
     .. warning::
 
         您不应当直接初始化该类，然后传入给 ``Trainer``，换句话说，您应当使用该类的子类 ``TorchSingleDriver`` 和 ``TorchDDPDriver``，而不是
-        该类本身；
+        该类本身。
 
     .. note::
 
-        您可以在使用 ``TorchSingleDriver`` 和 ``TorchDDPDriver`` 时使用 ``TorchDriver`` 提供的接口；
+        您可以在使用 ``TorchSingleDriver`` 和 ``TorchDDPDriver`` 时使用 ``TorchDriver`` 提供的接口。
 
-    :param model: 训练时使用的 **pytorch** 模型；
+    :param model: 训练时使用的 **pytorch** 模型。
     :param fp16: 是否开启混合精度训练;
     :param torch_kwargs:
     """
@@ -70,6 +70,9 @@ class TorchDriver(Driver):
         self.wo_auto_param_call = kwargs.get("model_wo_auto_param_call", False)
 
     def zero_grad(self):
+        """
+        实现梯度置零的过程
+        """
         for optimizer in self.optimizers:
             self._clear_grad(optimizer, self.set_grad_to_none)
 
@@ -88,14 +91,25 @@ class TorchDriver(Driver):
                         p.grad.zero_()
 
     def backward(self, loss):
+        """
+        对 ``loss`` 进行反向传播
+        """
         self.grad_scaler.scale(loss).backward()
 
     def step(self):
+        r"""
+        实现参数的优化更新过程
+        """
         for optimizer in self.optimizers:
             self.grad_scaler.step(optimizer)
             self.grad_scaler.update()
 
     def check_dataloader_legality(self, dataloader):
+        """
+        检测 DataLoader 是否合法。支持的类型包括 :class:`~fastNLP.core.dataloaders.TorchDataLoader`、 :class:`torch.utils.data.DataLoader` 。
+
+        :param dataloder:
+        """
         if not isinstance(dataloader, DataLoader) and not isinstance(dataloader, OverfitDataLoader):
             raise TypeError(f"{DataLoader} is expected, instead of `{type(dataloader)}`")
         if len(dataloader) == 0:
@@ -112,11 +126,11 @@ class TorchDriver(Driver):
     @staticmethod
     def tensor_to_numeric(tensor, reduce: str = None):
         r"""
-        将 ``torch.Tensor`` 转换成 python 中的数值类型；
+        将 ``torch.Tensor`` 转换成 python 中的数值类型。
 
-        :param tensor: ``torch.Tensor``；
-        :param reduce: 当 tensor 是一个多数值的张量时，应当使用何种归一化操作来转换成单一数值，应当为以下类型之一：``['max', 'min', 'sum', 'mean']``；
-        :return: 返回一个单一数值，其数值类型是 python 中的基本的数值类型，例如 ``int，float`` 等；
+        :param tensor: ``torch.Tensor``。
+        :param reduce: 当 tensor 是一个多数值的张量时，应当使用何种归一化操作来转换成单一数值，应当为以下类型之一：``['max', 'min', 'sum', 'mean']``。
+        :return: 一个单一数值，其数值类型是 python 中的基本的数值类型，例如 ``int，float`` 等。
         """
 
         if tensor is None:
@@ -137,8 +151,9 @@ class TorchDriver(Driver):
 
     def set_model_mode(self, mode: str):
         r"""
-        设置模型的状态是 ``train`` 还是 ``eval``；
-        :param mode: ``train`` 或者 ``eval``；
+        设置模型为 ``train`` 或 ``eval`` 的模式；目的是为切换模型的训练和推理（会关闭 dropout 等）模式。
+
+        :param mode: 应为二者之一：``["train", "eval"]``
         """
         assert mode in {"train", "eval"}
         getattr(self.model, mode)()
@@ -146,10 +161,10 @@ class TorchDriver(Driver):
     @rank_zero_call
     def save_model(self, filepath: Union[str, Path], only_state_dict: bool = True, **kwargs):
         """
-        保存当前 driver 的模型到 folder 下。
+        保存当前 driver 的模型到 ``filepath``。
 
-        :param filepath: 保存到哪个文件夹；
-        :param only_state_dict: 是否只保存权重；
+        :param filepath: 保存文件的文件位置
+        :param only_state_dict: 是否只保存权重
         :return:
         """
         model = self.unwrap_model()
@@ -169,12 +184,10 @@ class TorchDriver(Driver):
 
     def load_model(self, filepath: Union[Path, str], only_state_dict: bool = True, **kwargs):
         """
-        从 folder 中加载权重并赋值到当前 driver 的模型上。
+        加载模型的函数；将 ``filepath`` 中的模型加载并赋值给当前 ``model`` 。
 
-        :param filepath: 加载权重或模型的路径
-        :param load_state_dict: 保存的内容是否只是权重。
-        :param kwargs:
-        :return:
+        :param filepath: 保存文件的文件位置
+        :param load_state_dict: 保存的内容是否只是权重
         """
         model = self.unwrap_model()
         res = torch.load(filepath, map_location='cpu')
@@ -190,6 +203,17 @@ class TorchDriver(Driver):
 
     @rank_zero_call
     def save_checkpoint(self, folder: Path, states: Dict, dataloader, only_state_dict: bool = True, should_save_model: bool = True, **kwargs):
+        r"""
+        断点重训的保存函数，该函数会负责保存 **优化器** 、 **sampler** 和 **fp16** 的状态，以及 **模型** （若 ``should_save_model`` 为 ``True``）
+
+        :param folder: 保存断点重训的状态的文件夹；:meth:`save_checkpoint` 函数应该在该路径下面下面新增名为 ``FASTNLP_CHECKPOINT_FILENAME`` 与
+            ``FASTNLP_MODEL_FILENAME`` （如果 ``should_save_model`` 为 ``True`` ）的文件。把 model 相关的内容放入到 ``FASTNLP_MODEL_FILENAME`` 文件
+            中，将传入的 ``states`` 以及自身产生的其它状态一并保存在 ``FASTNLP_CHECKPOINT_FILENAME`` 里面。
+        :param states: 由 :class:`~fastNLP.core.controllers.Trainer` 传入的一个字典，其中已经包含了为了实现断点重训所需要保存的其它对象的状态。
+        :param dataloader: 正在使用的 dataloader。
+        :param only_state_dict: 是否只保存模型的参数，当 ``should_save_model`` 为 ``False`` ，该参数无效。
+        :param should_save_model: 是否应该保存模型，如果为 ``False`` ，Driver 将不负责 model 的保存。
+        """
         # 传入的 dataloader 参数是 trainer 的 dataloader 属性，因为 driver 的所有 dataloader 我们是不会去改变它的，而是通过改变
         #  trainer.dataloader 来改变 dataloader 的状态，从而适配训练或者评测环境；
 
@@ -297,6 +321,36 @@ class TorchDriver(Driver):
         logger.debug("Load optimizer state dict.")
 
     def load_checkpoint(self, folder: Path, dataloader, only_state_dict: bool = True, should_load_model: bool = True, **kwargs) -> Dict:
+        r"""
+        断点重训的加载函数，该函数会负责读取数据，并且恢复 **优化器** 、**sampler** 、 **fp16** 的状态和 **模型** （如果 ``should_load_model`` 为 True）以及其它
+        在 :meth:`save_checkpoint` 函数中执行的保存操作，然后将一个 state 字典返回给 :class:`~fastNLP.core.controllers.Trainer` （ 内容为 :meth:`save_checkpoint` 
+        接受到的 ``states`` ）。
+
+        该函数应该在所有 rank 上执行。
+
+        :param folder: 读取该 folder 下的 ``FASTNLP_CHECKPOINT_FILENAME`` 文件与 ``FASTNLP_MODEL_FILENAME``
+            （如果 should_load_model 为True）。
+        :param dataloader: 当前给定 dataloader，需要根据保存的 dataloader 状态合理设置。若该值为 ``None`` ，则不需要返回 ``'dataloader'``
+            以及 ``'batch_idx_in_epoch'`` 这两个值。
+        :param only_state_dict: 是否仅读取模型的 state_dict ，当 ``should_save_model`` 为 ``False`` ，该参数无效。如果为 ``True`` ，说明保存的内容为权重；如果为
+            False 说明保存的是模型，但也是通过当前 Driver 的模型去加载保存的模型的权重，而不是使用保存的模型替换当前模型。
+        :param should_load_model: 是否应该加载模型，如果为 ``False`` ，Driver 将不负责加载模型。若该参数为 ``True`` ，但在保存的状态中没有
+            找到对应的模型状态，则报错。
+        :return: :meth:`save_checkpoint` 函数输入的 ``states`` 内容。除此之外，还返回的内容有：
+
+            * *dataloader* -- 根据传入的 ``dataloader`` 与读取出的状态设置为合理状态的 dataloader。在当前 ``dataloader`` 样本数与读取出的 sampler 样本数
+              不一致时报错。
+            * *batch_idx_in_epoch* -- :class:`int` 类型的数据，表明当前 epoch 进行到了第几个 batch 。请注意，该值不能仅通过保存的数据中读取的，因为前后两次运行的
+              ``batch_size`` 可能有变化，而应该符合以下等式::
+
+                返回的 dataloader 还会产生的 batch 数量 + batch_idx_in_epoch = 原来不断点训练时的 batch 的总数
+              
+              由于 ``返回的 dataloader 还会产生的batch数`` 在 ``batch_size`` 与 ``drop_last`` 参数给定的情况下，无法改变，因此只能通过调整 ``batch_idx_in_epoch``
+              这个值来使等式成立。一个简单的计算原则如下：
+
+                * drop_last 为 ``True`` 时，等同于 floor(sample_in_this_rank/batch_size) - floor(num_left_samples/batch_size)；
+                * drop_last 为 ``False`` 时，等同于 ceil(sample_in_this_rank/batch_size) - ceil(num_left_samples/batch_size)。
+        """
         states = torch.load(folder.joinpath(FASTNLP_CHECKPOINT_FILENAME))
 
         # 1. 加载 optimizers 的状态；
@@ -326,29 +380,33 @@ class TorchDriver(Driver):
 
     def get_evaluate_context(self):
         r"""
-        :return: 返回 ``torch.no_grad`` 这个 context；
+        返回一个不计算梯度的上下文环境用来对模型进行评测。
+
+        :return: 上下文环境 ``torch.no_grad`` 
         """
         return torch.no_grad
 
     @staticmethod
     def move_model_to_device(model: "torch.nn.Module", device: "torch.device"):
         r"""
-        将模型迁移到对应的设备上；
+        将模型迁移到对应的设备上
         """
         if device is not None:
             model.to(device)
 
     def move_data_to_device(self, batch):
         """
-        将一个 batch 的数据迁移到对应的设备上；
+        将一个 ``batch`` 的数据迁移到对应的设备上
 
-        :param batch: 一个 batch 的数据，可以是 ``list、dict`` 等；
-        :return:
+        :param batch: 包含 :class:`torch.Tensor` 的数据集合，可以是 **List**、**Dict** 等嵌套类型
+        :return: 移动到指定机器后的 ``batch``
         """
         return torch_move_data_to_device(batch, self.data_device, self.non_blocking)
 
     @staticmethod
     def worker_init_function(worker_id: int, rank: Optional[int] = None) -> None:  # pragma: no cover
+        """
+        """
         """The worker_init_fn that Lightning automatically adds to your dataloader if you previously set the seed
         with ``seed_everything(seed, workers=True)``.
 
@@ -371,11 +429,20 @@ class TorchDriver(Driver):
         random.seed(stdlib_seed)
 
     def set_deterministic_dataloader(self, dataloader: "DataLoader"):
+        """
+        为了确定性训练要对 ``dataloader`` 进行修改，保证在确定随机数种子后，每次重新训练得到的结果是一样的。 
+        """
         if dataloader.worker_init_fn is None:
             dataloader.worker_init_fn = partial(self.worker_init_function,
                                                 rank=int(os.environ.get(FASTNLP_GLOBAL_RANK, 0)))
 
     def set_sampler_epoch(self, dataloader: "DataLoader", cur_epoch_idx: int):
+        r"""
+        对于分布式的 ``sampler``，需要在每一个 ``epoch`` 前设置随机数种子，来保证每一个进程上的 ``shuffle`` 是一样的。
+
+        :param dataloader: 需要设置 ``epoch`` 的 ``dataloader``
+        :param cur_epoch_idx: 当前是第几个 ``epoch``
+        """
         # 保证 ddp 训练时的 shuffle=True 时的正确性，因为需要保证每一个进程上的 sampler 的shuffle 的随机数种子是一样的；
         if callable(getattr(dataloader.sampler, "set_epoch", None)):
             dataloader.sampler.set_epoch(cur_epoch_idx)
@@ -383,9 +450,9 @@ class TorchDriver(Driver):
     @staticmethod
     def get_dataloader_args(dataloader: "DataLoader"):
         """
-        获取 dataloader 的 shuffle 和 drop_last 属性；
+        从 ``dataloader`` 中获取参数 ``dataset``, ``batch_sampler``, ``sampler``, ``batch_size``, ``shuffle`` 
+        和 ``drop_last`` 。
         """
-
         @dataclass
         class Res:
             dataset: Optional[Dataset] = None

fastNLP/core/drivers/torch_driver/torch_fsdp.py

@@ -0,0 +1,381 @@
+
+
+
+from fastNLP.envs.imports import _TORCH_GREATER_EQUAL_1_12
+
+if _TORCH_GREATER_EQUAL_1_12:
+    from torch.distributed.fsdp import FullyShardedDataParallel, StateDictType, FullStateDictConfig, OptimStateKeyType
+
+import os
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+from typing import Optional, Union, List, Dict, Mapping
+from pathlib import Path
+
+from .ddp import TorchDDPDriver
+from fastNLP.core.drivers.torch_driver.utils import (
+    _DDPWrappingModel,
+)
+
+from fastNLP.envs import FASTNLP_DISTRIBUTED_CHECK, FASTNLP_MODEL_FILENAME, FASTNLP_CHECKPOINT_FILENAME, \
+    FASTNLP_GLOBAL_RANK, rank_zero_call
+from fastNLP.core.drivers.torch_driver.utils import DummyGradScaler
+from fastNLP.core.log import logger
+from fastNLP.core.utils import check_user_specific_params
+from .utils import optimizer_state_to_device
+
+
+"""
+参考文档：
+1. https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/
+2. https://pytorch.org/docs/stable/fsdp.html?highlight=fsdp
+3. https://pytorch.org/tutorials/intermediate/FSDP_tutorial.html
+4. https://engineering.fb.com/2021/07/15/open-source/fsdp/
+"""
+
+class TorchFSDPDriver(TorchDDPDriver):
+    r"""
+    实现对于 pytorch 自己实现的 fully sharded data parallel；请阅读
+    `该文档 <https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.FullyShardedDataParallel.full_optim_state_dict>`_
+    了解更多：
+
+    .. note::
+
+        ``TorchFSDPDriver`` 大部分行为与 ``TorchDDPDriver`` 相同，如果您不了解 ``TorchDDPDriver``，
+        您可以先阅读 :class:`~fastNLP.core.drivers.TorchDDPDriver`；
+
+    .. warning::
+
+        ``TorchFSDPDriver`` 现在还不支持断点重训功能，但是支持保存模型和加载模型；
+
+        注意当您在加载和保存模型的 checkpointcallback 的时候，您可以通过在初始化 ``Trainer`` 时传入
+        ``torch_kwargs={"fsdp_kwargs": {'save_on_rank0': True/False, 'load_on_rank0': True/False}}`` 来指定保存模型的行为：
+
+            1. save/load_on_rank0 = True：表示在加载和保存模型时将所有 rank 上的模型参数全部聚合到 rank0 上，注意这样可能会造成 OOM；
+            2. save/load_on_rank0 = False：表示每个 rank 分别保存加载自己独有的模型参数；
+
+    :param model: 传入给 ``Trainer`` 的 ``model`` 参数
+    :param parallel_device: 用于分布式训练的 ``gpu`` 设备
+    :param is_pull_by_torch_run: 标志当前的脚本的启动是否由 ``python -m torch.distributed.launch`` 启动的
+    :param fp16: 是否开启 fp16 训练
+    :param torch_kwargs: 
+
+        * *fsdp_kwargs* --
+        * *set_grad_to_none* -- 是否在训练过程中在每一次 optimizer 更新后将 grad 置为 ``None``
+        * *non_blocking* -- 表示用于 :meth:`torch.Tensor.to` 方法的参数 non_blocking
+        * *gradscaler_kwargs* -- 用于 ``fp16=True`` 时，提供给 :class:`torch.amp.cuda.GradScaler` 的参数
+    :kwargs:
+        * *wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为
+
+        .. note::
+
+            关于该参数的详细说明，请参见 :class:`~fastNLP.core.controllers.Trainer` 中的描述；函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`。
+
+    """
+
+    def __init__(
+            self,
+            model,
+            parallel_device: Optional[Union[List["torch.device"], "torch.device"]],
+            is_pull_by_torch_run: bool = False,
+            fp16: bool = False,
+            torch_kwargs: Dict = None,
+            **kwargs
+    ):
+
+        # 在加入很多东西后，需要注意这里调用 super 函数的位置；
+        super(TorchDDPDriver, self).__init__(model, fp16=fp16, torch_kwargs=torch_kwargs, **kwargs)
+
+        if isinstance(model, torch.nn.DataParallel):
+            raise ValueError(f"Parameter `model` can not be `DataParallel` in `TorchDDPDriver`, it should be "
+                             f"`torch.nn.Module` or `torch.nn.parallel.DistributedDataParallel` type.")
+
+        # 如果用户自己在外面初始化 DDP，那么其一定是通过 python -m torch.distributed.launch 拉起的；
+        self.is_pull_by_torch_run = is_pull_by_torch_run
+        self.parallel_device = parallel_device
+        if not is_pull_by_torch_run and parallel_device is None:
+            raise ValueError(
+                "Parameter `parallel_device` can not be None when using `TorchDDPDriver`. This error is caused "
+                "when your value of parameter `device` is `None` in your `Trainer` instance.")
+
+        # 注意我们在 initialize_torch_driver 中的逻辑就是如果是 is_pull_by_torch_run，那么我们就直接把 parallel_device 置为当前进程的gpu；
+        if is_pull_by_torch_run:
+            self.model_device = parallel_device
+        else:
+            # 我们的 model_device 一定是 torch.device，而不是一个 list；
+            self.model_device = parallel_device[self.local_rank]
+
+        # 如果用户自己在外面初始化了 FSDP；
+        self.outside_ddp = False
+        if dist.is_initialized() and FASTNLP_DISTRIBUTED_CHECK not in os.environ and \
+                "fastnlp_torch_launch_not_ddp" not in os.environ:
+            # 如果用户自己在外面初始化了 DDP，那么我们要求用户传入的模型一定是已经由 DistributedDataParallel 包裹后的模型；
+            if not isinstance(model, FullyShardedDataParallel):
+                raise RuntimeError(
+                    "It is not allowed to input a normal model instead of `FullyShardedDataParallel` when"
+                    "you initialize the ddp process out of our control.")
+            if isinstance(model, DistributedDataParallel):
+                logger.warning("You are using `TorchFSDPDriver`, but you have initialized your model as "
+                               "`DistributedDataParallel`, which will make the `FullyShardedDataParallel` not work "
+                               "as expected. You could just delete `DistributedDataParallel` wrap operation.")
+
+            self.outside_ddp = True
+            # 用户只有将模型上传到对应机器上后才能用 DistributedDataParallel 包裹，因此如果用户在外面初始化了 DDP，那么在 TorchDDPDriver 中
+            #  我们就直接将 model_device 置为 None；
+            self.model_device = None
+
+        # 当用户自己在外面初始化 DDP 时我们会将 model_device 置为 None，这是用户可以通过 `data_device` 将对应的数据移到指定的机器上;
+        self._data_device = kwargs.get("data_device", None)
+        if isinstance(self._data_device, int):
+            if self._data_device < 0:
+                raise ValueError("Parameter `data_device` can not be smaller than 0.")
+            _could_use_device_num = torch.cuda.device_count()
+            if self._data_device >= _could_use_device_num:
+                raise ValueError("The gpu device that parameter `device` specifies is not existed.")
+            self._data_device = torch.device(f"cuda:{self._data_device}")
+        elif isinstance(self._data_device, str):
+            self._data_device = torch.device(self._data_device)
+        elif self._data_device is not None and not isinstance(self._data_device, torch.device):
+            raise ValueError("Parameter `device` is wrong type, please check our documentation for the right use.")
+
+        self._master_port = None
+        # world_size 表示的就是全局的显卡的数量；
+        self.world_size = None  # int(os.environ.get("WORLD_SIZE"))  len(self.parallel_device)
+        self.global_rank = 0
+
+        self._fsdp_kwargs = self._torch_kwargs.get("fsdp_kwargs", {})
+        self._save_on_rank0 = self._fsdp_kwargs.get("save_on_rank0", False)
+        if "save_on_rank0" in self._fsdp_kwargs:
+            self._fsdp_kwargs.pop("save_on_rank0")
+        self._load_on_rank0 = self._fsdp_kwargs.get("load_on_rank0", False)
+        if "load_on_rank0" in self._fsdp_kwargs:
+            self._fsdp_kwargs.pop("load_on_rank0")
+
+        if self._save_on_rank0 != self._load_on_rank0:
+            logger.warning(f"Notice the behavior between ``save`` and ``load`` is not matched, you choose "
+                           f"{'save on rank0' if self._save_on_rank0 else 'save on each rank'}, but "
+                           f"{'load on rank0' if self._save_on_rank0 else 'load on each rank'}!")
+
+        check_user_specific_params(self._fsdp_kwargs, FullyShardedDataParallel.__init__, FullyShardedDataParallel.__name__)
+        if "cpu_offload" in self._fsdp_kwargs and kwargs["accumulation_steps"] != 1:
+            logger.warning("It is not supported ``accumulation_steps`` when using ``cpu_offload`` in "
+                           "``FullyShardedDataParallel``.")
+
+        self.output_from_new_proc = kwargs.get("output_from_new_proc", "only_error")
+        assert isinstance(self.output_from_new_proc, str), "Parameter `output_from_new_proc` can only be `str` type."
+        if self.output_from_new_proc not in {"all", "ignore", "only_error"}:
+            os.makedirs(name=self.output_from_new_proc, exist_ok=True)
+            self.output_from_new_proc = os.path.abspath(self.output_from_new_proc)
+
+        self._has_setup = False  # 设置这一参数是因为 evaluator 中也会进行 setup 操作，但是显然是不需要的也不应该的；
+        self._has_ddpwrapped = False  # 判断传入的模型是否经过 _has_ddpwrapped 包裹；
+
+    def configure_ddp(self):
+        torch.cuda.set_device(self.model_device)
+        if not isinstance(self.model, FullyShardedDataParallel):
+            self.model = FullyShardedDataParallel(
+                # 注意这里的 self.model_device 是 `torch.device` type，因此 self.model_device.index；
+                _DDPWrappingModel(self.model), device_id=self.model_device.index,
+                **self._fsdp_kwargs
+            )
+
+            # 必须先使用 FullyShardedDataParallel 包裹模型后再使用 optimizer 包裹模型的参数，因此这里需要将 optimizer 重新初始化一遍；
+            for i in range(len(self.optimizers)):
+                self.optimizers[i] = type(self.optimizers[i])(self.model.parameters(), **self.optimizers[i].defaults)
+
+            self._has_ddpwrapped = True
+
+    def unwrap_model(self):
+        """
+        注意该函数因为需要在特定的时候进行调用，例如 ddp 在 get_model_call_fn 的时候，因此不能够删除；
+        如果您使用该函数来获取原模型的结构信息，是可以的；
+        但是如果您想要通过该函数来获取原模型实际的参数，是不可以的，因为在 FullyShardedDataParallel 中模型被切分成了多个部分，而对于每个 gpu 上
+        的模型只是整体模型的一部分。
+        """
+        _module = self.model.module.module
+        if isinstance(_module, _DDPWrappingModel):
+            return _module.model
+        else:
+            return _module
+
+    def save_model(self, filepath: Union[str, Path], only_state_dict: bool = True, **kwargs):
+        """
+        保存的模型到 ``filepath`` 中。
+
+        :param filepath: 文件路径
+        :param only_state_dict: 是否只保存权重；在 ``TorchFSDPDriver`` 中只能为 ``True`` 。
+        :param kwargs:
+        :return:
+        """
+        filepath = Path(filepath)
+        prefix = filepath.parent
+        filename = filepath.name
+        _filename = filename.split('.')
+        filename, suffix = _filename[0], '.'.join(_filename[1:])
+        if only_state_dict:
+            if self._save_on_rank0:
+                full_state_dict_config = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
+                with FullyShardedDataParallel.state_dict_type(self.model, StateDictType.FULL_STATE_DICT, full_state_dict_config):
+                    state_dict = self.model.state_dict()
+                rank_zero_call(torch.save)(state_dict, filepath)
+            else:
+                # 添加 'rank0/1' 字段来区分全部聚集到 rank0 保存的方式；
+                _filename = filename.split('_')
+                filename = _filename[0] + f"_rank{int(os.environ.get(FASTNLP_GLOBAL_RANK, 0))}_" + _filename[1]
+                filepath = prefix.joinpath(filename + "." + suffix)
+                with FullyShardedDataParallel.state_dict_type(self.model, StateDictType.LOCAL_STATE_DICT):
+                    state_dict = self.model.state_dict()
+                torch.save(state_dict, filepath)
+        else:
+            raise RuntimeError("When using `TorchFSDPDriver`, only `only_state_dict=True` is allowed.")
+
+    def load_model(self, filepath: Union[Path, str], only_state_dict: bool = True, **kwargs):
+        """
+        从 ``filepath`` 中加载权重并赋值到当前 driver 的模型上。
+
+        :param filepath: 加载权重或模型的路径
+        :param load_state_dict: 保存的内容是否只是权重；在 ``TorchFSDPDriver`` 中只能为 ``True`` 。
+        :param kwargs:
+        :return:
+        """
+        if only_state_dict is False:
+            raise RuntimeError("When using `TorchFSDPDriver`, only `only_state_dict=True` is allowed.")
+        filepath = Path(filepath)
+        prefix = filepath.parent
+        filename = filepath.name
+        _filename = filename.split('.')
+        filename, suffix = _filename[0], '.'.join(_filename[1:])
+
+        if not self._load_on_rank0:
+            _filename = filename.split('_')
+            filename = _filename[0] + f"_rank{int(os.environ.get(FASTNLP_GLOBAL_RANK, 0))}_" + _filename[1]
+            filepath = prefix.joinpath(filename + "." + suffix)
+            states = torch.load(filepath)
+        else:
+            states = torch.load(filepath, map_location="cpu")
+
+        if isinstance(states, dict) and only_state_dict is False:
+            logger.rank_zero_warning(f"It seems like that {filepath} only contains state, you may need to use "
+                                     f"`only_state_dict=True`")
+        elif not isinstance(states, dict) and only_state_dict is True:
+            logger.rank_zero_warning(f"It seems like that {filepath} is not state, you may need to use "
+                                     f"`only_state_dict=False`")
+        if not isinstance(states, Mapping):
+            states = states.state_dict()
+
+        if self._load_on_rank0:
+            with FullyShardedDataParallel.state_dict_type(self.model, StateDictType.FULL_STATE_DICT):
+                self.model.load_state_dict(states)
+        else:
+            with FullyShardedDataParallel.state_dict_type(self.model, StateDictType.LOCAL_STATE_DICT):
+                self.model.load_state_dict(states)
+
+    def save_checkpoint(self, folder: Path, states: Dict, dataloader, only_state_dict: bool = True, should_save_model: bool = True, **kwargs):
+        raise RuntimeError("``TorchFSDPDriver`` does not support ``save_checkpoint`` function for now, there is some "
+                           "technical issues that needs to solve. You can implement your own breakpoint retraining "
+                           "by rewriting this function. The important thing is how to save and load the optimizers' state dict, "
+                           "you can see ``https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.FullyShardedDataParallel.full_optim_state_dict``.")
+
+    def load_checkpoint(self, folder: Path, dataloader, only_state_dict: bool = True, should_load_model: bool = True, **kwargs) -> Dict:
+        raise RuntimeError("``TorchFSDPDriver`` does not support ``load_checkpoint`` function for now, there is some "
+                           "technical issues that needs to solve. You can implement your own breakpoint retraining "
+                           "by rewriting this function. The important thing is how to save and load the optimizers' state dict, "
+                           "you can see ``https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.FullyShardedDataParallel.full_optim_state_dict``.")
+
+    # todo 这些加了 __ 的函数是目前还不支持；
+    #  这是因为 1.12 的 pytorch fsdp 的关于如何保存和加载 optimizer state dict 的接口有点过于反人类，无法在 fastNLP 的框架中进行调和
+    #  使用；
+    def __get_optimizer_state(self):
+        optimizers_state_dict = {}
+        for i in range(len(self.optimizers)):
+            # 注意这里其余 rank 拿到的是一个空字典，因此在真正保存的时候需要保证只有 rank0 在工作；
+            optimizer_state = FullyShardedDataParallel.full_optim_state_dict(self.model, self.optimizers[i])
+            if self._save_on_rank0:
+                with FullyShardedDataParallel.summon_full_params(self.model):
+                    if int(os.environ.get(FASTNLP_GLOBAL_RANK, 0)) == 0:
+                        unwrapped_model = self.model.module.module
+                        optimizer_state = FullyShardedDataParallel.rekey_optim_state_dict(
+                            optimizer_state, OptimStateKeyType.PARAM_ID, unwrapped_model)
+                if int(os.environ.get(FASTNLP_GLOBAL_RANK, 0)) == 0:
+                    optimizer_state["state"] = optimizer_state_to_device(optimizer_state["state"], torch.device("cpu"))
+            optimizers_state_dict[f"optimizer{i}"] = optimizer_state  # 注意这里没有使用 deepcopy，测试是不需要的；
+        return optimizers_state_dict
+
+    # 这里单独拿出来是因为对于 fsdp 来说，每一个进程都需要运行此函数，因此不能包裹 rank_zero_call；
+    def __save_checkpoint(self, folder: Path, states: Dict, dataloader, only_state_dict: bool = True, should_save_model: bool = True, **kwargs):
+        if not only_state_dict:
+            raise RuntimeError("When using `TorchFSDPDriver`, only `only_state_dict=True` is allowed.")
+
+        # 1. sampler 的状态；
+        num_consumed_batches = states.pop('num_consumed_batches')
+        states['sampler_states'] = self.get_sampler_state(dataloader, num_consumed_batches)
+
+        # 2. 保存模型的状态；
+        if should_save_model:
+            if not os.path.exists(folder):
+                os.mkdir(folder)
+            model_path = folder.joinpath(FASTNLP_MODEL_FILENAME)
+            self.save_model(model_path, only_state_dict=True)
+
+        # 3. 保存 optimizers 的状态；
+        states["optimizers_state_dict"] = self.get_optimizer_state()
+        logger.debug("Save optimizer state dict.")
+
+        # 4. 保存fp16的状态
+        if not isinstance(self.grad_scaler, DummyGradScaler):
+            grad_scaler_state_dict = self.grad_scaler.state_dict()
+            states['grad_scaler_state_dict'] = grad_scaler_state_dict
+
+        # 确保只有 rank0 才会执行实际的保存操作；
+        rank_zero_call(torch.save)(states, Path(folder).joinpath(FASTNLP_CHECKPOINT_FILENAME))
+
+    def __load_optimizer_state(self, states):
+        assert len(states) == len(self.optimizers), f"The number of optimizers is:{len(self.optimizers)}, while in " \
+                                                    f"checkpoint it is:{len(states)}"
+
+        with FullyShardedDataParallel.summon_full_params(self.model):
+            unwrapped_model = self.model.module.module
+
+            for i in range(len(self.optimizers)):
+                optimizer_state = states[f'optimizer{i}']
+                if self._load_on_rank0:
+                    optimizer_state = FullyShardedDataParallel.rekey_optim_state_dict(optimizer_state, OptimStateKeyType.PARAM_NAME, unwrapped_model)
+                optimizer_state = FullyShardedDataParallel.shard_full_optim_state_dict(optimizer_state, unwrapped_model)
+                optimizer: torch.optim.Optimizer = type(self.optimizers[i])(unwrapped_model.parameters(), **self.optimizers[i].defaults)
+                optimizer.load_state_dict(optimizer_state)
+                self.optimizers[i] = optimizer
+
+        logger.debug("Load optimizer state dict.")
+
+    def __load_checkpoint(self, folder: Path, dataloader, only_state_dict: bool = True, should_load_model: bool = True, **kwargs) -> Dict:
+        if not only_state_dict:
+            raise RuntimeError("When using `TorchFSDPDriver`, only `only_state_dict=True` is allowed.")
+
+        states = torch.load(folder.joinpath(FASTNLP_CHECKPOINT_FILENAME))
+
+        # 1. 加载 optimizers 的状态；
+        optimizers_state_dict = states.pop("optimizers_state_dict")
+        self.load_optimizer_state(optimizers_state_dict)
+
+        # 2. 加载模型状态；
+        if should_load_model:
+            self.load_model(filepath=folder.joinpath(FASTNLP_MODEL_FILENAME), only_state_dict=only_state_dict)
+
+        # 3. 加载 fp16 的状态
+        if "grad_scaler_state_dict" in states:
+            grad_scaler_state_dict = states.pop("grad_scaler_state_dict")
+            if not isinstance(self.grad_scaler, DummyGradScaler):
+                self.grad_scaler.load_state_dict(grad_scaler_state_dict)
+                logger.debug("Load grad_scaler state dict...")
+        elif not isinstance(self.grad_scaler, DummyGradScaler):
+            logger.rank_zero_warning(f"Checkpoint {folder} is not trained with fp16=True, while resume to a fp16=True training, "
+                           f"the training process may be unstable.")
+
+        # 4. 恢复 sampler 的状态；
+        sampler_states = states.pop('sampler_states')
+        states_ret = self.load_sampler_state(dataloader, sampler_states)
+        states.update(states_ret)
+
+        return states
+

fastNLP/core/drivers/torch_driver/utils.py

@@ -42,7 +42,7 @@ def torch_seed_everything(seed: int = None, add_global_rank_to_seed: bool = True
 
     :param seed: 全局随机状态的整数值种子。如果为 ``None`` 则会根据时间戳生成一个种子。
     :param add_global_rank_to_seed: 在分布式训练中，是否在不同 **rank** 中使用不同的随机数。
-        当设置为 ``True`` 时，**FastNLP** 会将种子加上当前的 ``global_rank``。
+        当设置为 ``True`` 时，**fastNLP** 会将种子加上当前的 ``global_rank``。
     """
     max_seed_value = np.iinfo(np.uint32).max
     min_seed_value = np.iinfo(np.uint32).min
@@ -290,11 +290,11 @@ def replace_batch_sampler(dataloader, new_batch_sampler):
 
 def optimizer_state_to_device(state, device):
     r"""
-    将一个 ``optimizer`` 的 ``state_dict`` 迁移到对应的设备；
+    将一个 ``optimizer`` 的 ``state_dict`` 迁移到对应的设备。
 
-    :param state: ``optimzier.state_dict()``；
-    :param device: 要迁移到的目的设备；
-    :return: 返回迁移后的新的 state_dict；
+    :param state: ``optimzier.state_dict()``。
+    :param device: 要迁移到的目的设备。
+    :return: 迁移后的新的 state_dict。
     """
     new_state = {}
     for name, param in state.items():

fastNLP/core/drivers/utils.py

@@ -1,18 +1,22 @@
 from typing import List
 import subprocess
 
+__all__ = []
 
 def distributed_open_proc(output_from_new_proc:str, command:List[str], env_copy:dict, rank:int=None):
     r"""
     使用 command 通过 subprocess.Popen 开启新的进程。
 
-    :param output_from_new_proc: 可选 ["ignore", "all", "only_error"]，以上三个为特殊关键字，分别表示完全忽略拉起进程的打印输出，
-        only_error 表示只打印错误输出流；all 表示子进程的所有输出都打印。如果不为以上的关键字，则表示一个文件夹，将在该文件夹下建立
-        两个文件，名称分别为 {rank}_std.log, {rank}_err.log 。原有的文件会被直接覆盖。
-    :param command: List[str] 启动的命令
+    :param output_from_new_proc: 可选 ``["ignore", "all", "only_error"]``，以上三个为特殊关键字，分别表示：
+        * ``"ignore:`` -- 完全忽略拉起进程的打印输出；
+        * ``"only_error"`` -- 表示只打印错误输出流；
+        * ``"all"`` -- 子进程的所有输出都打印。
+        * 如果不为以上的关键字，则表示一个文件夹，将在该文件夹下建立两个文件，名称分别为 {rank}_std.log, {rank}_err.log 。
+          原有的文件会被直接覆盖。
+    :param command: 启动的命令
     :param env_copy: 需要注入的环境变量。
     :param rank: global_rank；
-    :return: 返回使用 ``subprocess.Popen`` 打开的进程；
+    :return: 使用 ``subprocess.Popen`` 打开的进程；
     """
     if output_from_new_proc == "all":
         proc = subprocess.Popen(command, env=env_copy)

fastNLP/core/log/handler.py

@@ -7,6 +7,7 @@ try:
 except ImportError:
     tqdm = None
 
+__all__ = []
 
 if tqdm is not None:
     class TqdmLoggingHandler(logging.Handler):

fastNLP/core/log/highlighter.py

@@ -1,6 +1,6 @@
 from rich.highlighter import Highlighter
 
-
+__all__ = []
 class ColorHighlighter(Highlighter):
     def __init__(self, color='black'):
         self.color = color

fastNLP/core/log/logger.py

@@ -1,20 +1,24 @@
 r"""
-Logger 是fastNLP中记录日志的模块，logger封装了logging模块的Logger，
-具体使用方式与直接使用logging.Logger相同，同时也新增一些简单好用的API
+:class:`Logger` 是 **fastNLP** 中记录日志的模块，**logger** 封装了 logging 模块的 Logger，
+具体使用方式与直接使用 :class:`logging.Logger` 相同，同时也新增一些简单好用的API
 
 使用方式::
 
-    from fastNLP import _logger
-    #
-    # _logger 可以和 logging.Logger 一样使用
-    _logger.info('your msg')
-    _logger.error('your msg')
+    from fastNLP import logger
 
-    # _logger 新增的API
+    # logger 可以和 logging.Logger 一样使用
+    logger.info('your msg')
+    logger.error('your msg')
+
+    # logger 新增的API
     # 将日志输出到文件，以及输出的日志等级
-    _logger.add_file('/path/to/log', level='INFO')
+    logger.add_file('/path/to/log', level='INFO')
     # 定义在命令行中的显示格式和日志等级
-    _logger.set_stdout('tqdm', level='WARN')
+    logger.set_stdout('tqdm', level='WARN')
+    # 仅警告一次
+    logger.warning_once('your msg')
+    # 分布式训练下，仅在 rank 0 输出警告
+    logger.rank_zero_warning('your msg')
 
 """
 

fastNLP/core/log/print.py

@@ -16,7 +16,7 @@ def print(*args, sep=' ', end='\n', file=None, flush=False):
 
     :param args: 需要打印的内容
     :param sep: 存在多个输入时，使用的间隔。
-    :param end: 该参数在当前设置无意义，因为结尾一定会被加入 '\\\\n' 。
+    :param end: 该参数在当前设置无意义，因为结尾一定会被加入 ``'\\\\n'`` 。
     :param file: 该参数无意义。
     :param flush: 该参数无意义。
     :return:

fastNLP/core/metrics/accuracy.py

@@ -14,24 +14,25 @@ from fastNLP.core.log import logger
 
 
 class Accuracy(Metric):
-    def __init__(self, backend: Union[str, Backend, None] = 'auto', aggregate_when_get_metric: bool = None):
-        """
-        计算 准确率 的 metric 。
+    """
+    计算 准确率 的 metric 。
 
-        :param backend: 目前支持四种类型的backend, ['auto', 'torch', 'paddle', 'jittor']。其中 auto 表示根据实际调用 Metric.update()
-            函数时传入的参数决定具体的 backend ，一般情况下直接使用 'auto' 即可。
-        :param aggregate_when_get_metric: 在计算 metric 的时候是否自动将各个进程上的相同的 element 的数字聚合后再得到 metric，
-            当 backend 不支持分布式时，该参数无意义。如果为 None ，将在 Evaluator 中根据 sampler 是否使用分布式进行自动设置。
-        """
+    :param backend: 目前支持五种类型的backend, ``['auto', 'torch', 'paddle', 'jittor', 'oneflow']``。其中 ``'auto'`` 表示根据实际调用
+        :meth:`update` 函数时传入的参数决定具体的 backend ，一般情况下直接使用 ``'auto'`` 即可。
+    :param aggregate_when_get_metric: 在计算 metric 的时候是否自动将各个进程上的相同的 element 的数字聚合后再得到 metric，
+        当 ``backend`` 不支持分布式时，该参数无意义。如果为 ``None`` ，将在 :class:`~fastNLP.core.controllers.Evaluator` 
+        中根据 ``sampler`` 是否使用分布式进行自动设置。
+    """
+    def __init__(self, backend: Union[str, Backend, None] = 'auto', aggregate_when_get_metric: bool = None):
         super(Accuracy, self).__init__(backend=backend, aggregate_when_get_metric=aggregate_when_get_metric)
         self.register_element(name='correct', value=0, aggregate_method='sum', backend=backend)
         self.register_element(name='total', value=0, aggregate_method="sum", backend=backend)
 
     def get_metric(self) -> dict:
         r"""
-        get_metric 函数将根据 update 函数累计的评价指标统计量来计算最终的评价结果.
+        :meth:`get_metric` 函数将根据 :meth:`update` 函数累计的评价指标统计量来计算最终的评价结果。
 
-        :return dict evaluate_result: {"acc": float, 'total': float, 'correct': float}
+        :return: 包含以下内容的字典：``{"acc": float, 'total': float, 'correct': float}``；
         """
         evaluate_result = {'acc': round(self.correct.get_scalar() / (self.total.get_scalar() + 1e-12), 6),
                            'total': self.total.item(), 'correct': self.correct.item()}
@@ -39,14 +40,14 @@ class Accuracy(Metric):
 
     def update(self, pred, target, seq_len=None):
         r"""
-        update 函数将针对一个批次的预测结果做评价指标的累计
-
-        :param pred: 预测的tensor, tensor的形状可以是torch.Size([B,]), torch.Size([B, n_classes]),
-                torch.Size([B, max_len]), 或者torch.Size([B, max_len, n_classes])
-        :param target: 真实值的tensor, tensor的形状可以是Element's can be: torch.Size([B,]),
-                torch.Size([B,]), torch.Size([B, max_len]), 或者torch.Size([B, max_len])
-        :param seq_len: 序列长度标记, 标记的形状可以是None, None, torch.Size([B]), 或者torch.Size([B]).
-                如果mask也被传进来的话seq_len会被忽略.
+        :meth:`update` 函数将针对一个批次的预测结果做评价指标的累计。
+
+        :param pred: 预测的 tensor, tensor 的形状可以是 ``[B,]`` 、``[B, n_classes]`` 、
+                ``[B, max_len]`` 或 ``[B, max_len, n_classes]``
+        :param target: 真实值的 tensor, tensor 的形状可以是 ``[B,]`` 、``[B,]`` 、``[B, max_len]``
+            或 ``[B, max_len]``
+        :param seq_len: 序列长度标记, 标记的形状可以是 ``None``,  或者 ``[B]`` 。
+                如果 mask 也被传进来的话 ``seq_len`` 会被忽略
         """
         # 为了兼容不同框架，我们将输入变量全部转为numpy类型来进行计算。
         pred = self.tensor2numpy(pred)
@@ -85,12 +86,11 @@ class Accuracy(Metric):
 
 class TransformersAccuracy(Accuracy):
     """
-    适配 transformers 中相关模型的 Accuracy metric 。
-
+    适配 :mod:`transformers` 中相关模型的 Accuracy metric 。
     """
     def update(self, logits, labels, attention_mask=None):
         r"""
-        update 函数将针对一个批次的预测结果做评价指标的累计
+        :meth:`update` 函数将针对一个批次的预测结果做评价指标的累计。
 
         :param logits: 形状为 ``[B, n_classes]`` 或 ``[B, max_len, n_classes]`` 。
         :param labels: 形状为 ``[B, ]`` 或 ``[B, max_len]``

fastNLP/core/metrics/backend/__init__.py

@@ -2,11 +2,15 @@ __all__ = [
     'Backend',
     'AutoBackend',
     'TorchBackend',
-    'PaddleBackend'
+    'PaddleBackend',
+    'JittorBackend',
+    'OneflowBackend',
 ]
 
 
 from .backend import Backend
 from .auto_backend import AutoBackend
-from .torch_backend.backend import TorchBackend
-from .paddle_backend.backend import PaddleBackend
+from .torch_backend import TorchBackend
+from .paddle_backend import PaddleBackend
+from .jittor_backend import JittorBackend
+from .oneflow_backend import OneflowBackend

fastNLP/core/metrics/backend/auto_backend.py

@@ -10,11 +10,11 @@ from .paddle_backend.backend import PaddleBackend
 from .jittor_backend.backend import JittorBackend
 from .oneflow_backend.backend import OneflowBackend
 
+__all__ = []
 
 class AutoBackend(Backend):
     """
-    不需要初始化 backend 的 AutoBackend,能够根据 get_metric 时候判断输入数据类型来选择 backend 是什么类型的
-
+    不需要初始化 ``backend`` 的 :class:`AutoBackend`，能够根据 :meth:`get_metric` 时候判断输入数据类型来选择 ``backend``。
     """
 
     def __init__(self, backend: Union[str, Backend, None]):
@@ -23,9 +23,9 @@ class AutoBackend(Backend):
 
         :param backend: 目前支持三种值，为 ``[str, Backend, None]``。
 
-            * 当 backend 为 `str` 时， 其只能为 'auto'
-            * 当 backend 为 ``Backend`` 对象时， 其直接使用该对象方法覆盖 AutoBackend
-            * 当 backend 为 ``None`` 时， 根据 get_metric 时候判断输入数据类型来选择 backend 是什么类型的
+            * 当 backend 为 :class:`str` 时， 其只能为 ``'auto'``；
+            * 当 backend 为 ``Backend`` 对象时， 其直接使用该对象方法覆盖 :class:`AutoBackend`；
+            * 当 backend 为 ``None`` 时， 根据 :meth:`get_metric` 时候判断输入数据类型来选择 ``backend``；
 
         """
         super(AutoBackend, self).__init__()
@@ -38,9 +38,10 @@ class AutoBackend(Backend):
 
         :param backend: 传入的 backend 值。
 
-            * 当 backend 为 `torch` 时， 选择 :class:`~fastNLP.core.metric.TorchBackend`
-            * 当 backend 为 `paddle` 时， 选择 :class:`~fastNLP.core.metric.PaddleBackend`
-            * 当 backend 为 `jittor` 时， 选择 :class:`~fastNLP.core.metric.JittorBackend`
+            * 当 backend 为 ``'torch'`` 时， 选择 :class:`~fastNLP.core.metric.TorchBackend`
+            * 当 backend 为 ``'paddle'` 时， 选择 :class:`~fastNLP.core.metric.PaddleBackend`
+            * 当 backend 为 ``'jittor'`` 时， 选择 :class:`~fastNLP.core.metric.JittorBackend`
+            * 当 backend 为 ``'oneflow'`` 时， 选择 :class:`~fastNLP.core.metric.OneflowBackend`
             * 当 backend 为 ``None`` 时， 直接初始化
 
         """
@@ -66,7 +67,7 @@ class AutoBackend(Backend):
         """
         根据 args 参数类型来选择需要真正初始化的 backend
 
-        :param args: args 参数， 可能为 ``jittor``, ``torch``, ``paddle``, ``numpy`` 类型， 能够检测并选择真正的 backend。
+        :param args: args 参数， 可能为 ``'jittor'``, ``'torch'``, ``'paddle'``, ``'oneflow'``, ``'numpy'`` 类型， 能够检测并选择真正的 backend。
 
         """
         assert not self.is_specified(), "This method should not be called after backend has been specified. " \

fastNLP/core/metrics/backend/backend.py

@@ -1,10 +1,10 @@
 from ..utils import AggregateMethodError
 
+__all__ = []
 
 class Backend:
     """
-    Backend 及其子类的所有方法都必须是无状态的。
-
+    执行评测时使用的 backend，是所有 backend 的父类。Backend 及其子类的所有方法都必须是无状态的。
     """
 
     def __init__(self):
@@ -12,7 +12,8 @@ class Backend:
 
     def aggregate(self, tensor, method: str):
         """
-        聚集结果，并根据 method 计算后，返回结果
+        聚集结果，并根据 ``method 计算后`` ，返回结果。
+
         :param tensor: 传入的张量
         :param method: 聚合的方法
         """
@@ -23,15 +24,15 @@ class Backend:
 
     def create_tensor(self, value: float):
         """
-        创建tensor，并且填入value作为值
+        创建 tensor，并且填入 ``value`` 作为值。
 
-        :param value: 需要初始化的 value 值
+        :param value: 需要初始化的 ``value`` 值
         """
         return value
 
     def fill_value(self, tensor, value: float):
         """
-        将tensor的值设置为value
+        将 tensor 的值设置为 ``value``
 
         :param tensor: 传进来的张量
         :param value: 需要填充的值
@@ -40,16 +41,16 @@ class Backend:
 
     def get_scalar(self, tensor) -> float:
         """
-        tensor的saclar值
+        ``tensor`` 的 saclar 值.
 
-        :param tensor: 传入的张量
+        :param tensor: 传入的张量;
         :return:
         """
         return tensor
 
     def is_specified(self) -> bool:
         """
-        判断是否是某种框架的 backend
+        判断是否是某种框架的 backend。
 
         :return:
         """
@@ -57,7 +58,7 @@ class Backend:
 
     def tensor2numpy(self, tensor):
         """
-        将 tensor 转为 numpy
+        将 ``tensor`` 转为 :class:`numpy.array`。
 
         :param tensor: 传入的张量
         :return:
@@ -66,16 +67,16 @@ class Backend:
 
     def move_tensor_to_device(self, tensor, device):
         """
-        将张量移动到某个设备上
+        将张量移动到某个设备上。
 
         :param tensor: 传入的张量
-        :param device: 设备号， 一般为 ``'cpu'``, ``'cuda:0'`` 等。
+        :param device: 设备号， 一般为 ``'cpu'``, ``'cuda:0'`` 等
         """
         return tensor
 
     def all_gather_object(self, obj, group=None):
         """
-        给定 obj 将各个 rank 上的 obj 汇总到每个 obj 上。返回一个 list 对象，里面依次为各个 rank 对应的 obj 。
+        给定 ``obj`` 将各个 rank 上的 ``obj`` 汇总到每个 ``obj`` 上。返回一个 :class:`list` 对象，里面依次为各个 rank 对应的 ``obj`` 。
 
         :param obj:
         :param group:

fastNLP/core/metrics/backend/jittor_backend/__init__.py

@@ -1 +1,5 @@
+__all__ = [
+    "JittorBackend",
+]
 
+from .backend import JittorBackend

fastNLP/core/metrics/backend/oneflow_backend/__init__.py

@@ -0,0 +1,5 @@
+__all__ = [
+    "OneflowBackend",
+]
+
+from .backend import OneflowBackend

fastNLP/core/metrics/backend/oneflow_backend/backend.py

@@ -25,12 +25,12 @@ class OneflowBackend(Backend):
         聚集结果，并根据 method 计算后，返回结果
 
         :param tensor: 需要聚合的张量
-        :param method: 聚合的方法， 目前支持 ``['sum', 'mean', 'max', 'mix']``:
+        :param method: 聚合的方法， 目前支持 ``['sum', 'mean', 'max', 'min']``:
 
             * method 为 ``'sum'`` 时， 会将多张卡上聚合结果在维度为 `0` 上 累加起来。
             * method 为 ``'mean'`` 时，会将多张卡上聚合结果在维度为 `0` 上取平均值。
             * method 为 ``'max'`` 时，会将多张卡上聚合结果在维度为 `0` 上取最大值。
-            * method 为 ``'mix'`` 时，会将多张卡上聚合结果在维度为 `0` 上取最小值。
+            * method 为 ``'min'`` 时，会将多张卡上聚合结果在维度为 `0` 上取最小值。
 
         """
         if isinstance(tensor, oneflow.Tensor):

fastNLP/core/metrics/backend/paddle_backend/__init__.py

@@ -2,4 +2,4 @@ __all__ = [
     'PaddleBackend'
 ]
 
-from .backend import Backend as PaddleBackend
+from .backend import PaddleBackend

fastNLP/core/metrics/backend/paddle_backend/backend.py

@@ -26,12 +26,12 @@ class PaddleBackend(Backend):
         聚集结果，并根据 method 计算后，返回结果
 
         :param tensor: 需要聚合的张量
-        :param method: 聚合的方法， 目前支持 ``['sum', 'mean', 'max', 'mix']``:
+        :param method: 聚合的方法， 目前支持 ``['sum', 'mean', 'max', 'min']``:
 
             * method 为 ``'sum'`` 时， 会将多张卡上聚合结果在维度为 `0` 上 累加起来。
             * method 为 ``'mean'`` 时，会将多张卡上聚合结果在维度为 `0` 上取平均值。
             * method 为 ``'max'`` 时，会将多张卡上聚合结果在维度为 `0` 上取最大值。
-            * method 为 ``'mix'`` 时，会将多张卡上聚合结果在维度为 `0` 上取最小值。
+            * method 为 ``'min'`` 时，会将多张卡上聚合结果在维度为 `0` 上取最小值。
 
         """
         if isinstance(tensor, paddle.Tensor):

fastNLP/core/metrics/backend/torch_backend/__init__.py

@@ -3,4 +3,4 @@ __all__ = [
 ]
 
 
-from .backend import Backend as TorchBackend
+from .backend import TorchBackend

fastNLP/core/metrics/backend/torch_backend/backend.py

@@ -24,12 +24,12 @@ class TorchBackend(Backend):
         聚集结果，并根据 method 计算后，返回结果
 
         :param tensor: 需要聚合的张量
-        :param method: 聚合的方法， 目前支持 ``['sum', 'mean', 'max', 'mix']``:
+        :param method: 聚合的方法， 目前支持 ``['sum', 'mean', 'max', 'min']``:
 
             * method 为 ``'sum'`` 时， 会将多张卡上聚合结果在维度为 `0` 上 累加起来。
             * method 为 ``'mean'`` 时，会将多张卡上聚合结果在维度为 `0` 上取平均值。
             * method 为 ``'max'`` 时，会将多张卡上聚合结果在维度为 `0` 上取最大值。
-            * method 为 ``'mix'`` 时，会将多张卡上聚合结果在维度为 `0` 上取最小值。
+            * method 为 ``'min'`` 时，会将多张卡上聚合结果在维度为 `0` 上取最小值。
 
         """
         if isinstance(tensor, torch.Tensor):

fastNLP/core/metrics/classify_f1_pre_rec_metric.py

@@ -14,27 +14,36 @@ from .utils import _compute_f_pre_rec
 from fastNLP.core.log import logger
 
 class ClassifyFPreRecMetric(Metric):
-    def __init__(self, tag_vocab: Vocabulary = None, ignore_labels: List[str] = None,
-                 only_gross: bool = True, f_type='micro', beta=1, backend: Union[str, Backend, None] = 'auto',
-                 aggregate_when_get_metric: bool = None) -> None:
-        """
+    """
+    计算分类结果 **F值** 的 **Metric** 。
 
-        :param tag_vocab: 标签的 :class:`~fastNLP.Vocabulary` . 默认值为 ``None``。若为 ``None`` 则使用数字来作为标签内容，
-        否则使用 vocab 来作为标签内容。
-        :param ignore_labels: ``str`` 组成的 ``list``. 这个 ``list``中的 class 不会被用于计算。例如在 POS tagging 时传入 ``['NN']``，
+    :param tag_vocab: 标签的 :class:`~fastNLP.core.Vocabulary` 。 默认值为 ``None``。若为 ``None`` 则使用数字来作为标签内容，
+        否则使用 vocab 来作为标签内容
+    :param ignore_labels: :class:`str` 组成的 :class:`list`. 这个 :class:`list` 中的 class 不会被用于计算。例如在 POS tagging 时传入 ``['NN']``，
         则不会计算 'NN' 个 label
-        :param only_gross: 是否只计算总的 ``f1``, ``precision``, ``recall``的值；如果为 ``False``，不仅返回总的 ``f1``, ``pre``,
+    :param only_gross: 是否只计算总的 ``f1``, ``precision``, ``recall``的值；如果为 ``False``，不仅返回总的 ``f1``, ``pre``,
         ``rec``, 还会返回每个 label 的 ``f1``, ``pre``, ``rec``
-        :param f_type: `micro` 或 `macro` .
-            * `micro` : 通过先计算总体的 TP，FN 和 FP 的数量，再计算 f, precision, recall;
-            * `macro` : 分布计算每个类别的 f, precision, recall，然后做平均（各类别 f 的权重相同）
-        :param beta: f_beta分数， :math:`f_{beta} = \frac{(1 + {beta}^{2})*(pre*rec)}{({beta}^{2}*pre + rec)}` .
-        :param backend: 目前支持四种类型的 backend, ``[torch, paddle, jittor, 'auto']``。其中 ``'auto'`` 表示根据实际调用 Metric.update()
-        函数时传入的参数决定具体的 backend ，大部分情况下直接使用 ``'auto'`` 即可。
-        :param aggregate_when_get_metric: 在计算 metric 的时候是否自动将各个进程上的相同的 element 的数字聚合后再得到metric，
-            当 backend 不支持分布式时，该参数无意义。如果为 ``None`` ，将在 Evaluator 中根据 sampler 是否使用分布式进行自动设置。
+    :param f_type: `micro` 或 `macro` 。
 
-        """
+        * `micro` : 通过先计算总体的 TP，FN 和 FP 的数量，再计算 f, precision, recall;
+        * `macro` : 分布计算每个类别的 f, precision, recall，然后做平均（各类别 f 的权重相同）
+
+    :param beta: **f_beta** 分数中的 ``beta`` 值。 常用为 ``beta=0.5, 1, 2`` 若为 0.5 则 **精确率** 的权重高于 **召回率** ；若为1，则两者平等；若为2，则
+        **召回率** 权重高于 **精确率** 。**f_beta** 分数的计算公式为：
+        
+        .. math::
+
+            f_{beta} = \\frac{(1 + {beta}^{2})*(pre*rec)}{({beta}^{2}*pre + rec)}
+
+    :param backend: 目前支持五种类型的 backend, ``['torch', 'paddle', 'jittor', 'oneflow', 'auto']``。其中 ``'auto'`` 表示根据实际调用 :meth:`update`
+        函数时传入的参数决定具体的 backend ，大部分情况下直接使用 ``'auto'`` 即可。
+    :param aggregate_when_get_metric: 在计算 metric 的时候是否自动将各个进程上的相同的 element 的数字聚合后再得到 metric，
+        当 backend 不支持分布式时，该参数无意义。如果为 ``None`` ，将在 :class:`~fastNLP.core.controllers.Evaluator` 中根据
+        sampler 是否使用分布式进行自动设置。
+    """
+    def __init__(self, tag_vocab: Vocabulary = None, ignore_labels: List[str] = None,
+                 only_gross: bool = True, f_type='micro', beta=1, backend: Union[str, Backend, None] = 'auto',
+                 aggregate_when_get_metric: bool = None) -> None:
         super(ClassifyFPreRecMetric, self).__init__(backend=backend,
                                                     aggregate_when_get_metric=aggregate_when_get_metric)
         if f_type not in ('micro', 'macro'):
@@ -56,8 +65,7 @@ class ClassifyFPreRecMetric(Metric):
 
     def reset(self):
         """
-        重置 tp, fp, fn 的值
-
+        重置 ``tp``, ``fp``, ``fn`` 的值
         """
         # 由于不是 element 了，需要自己手动清零一下
         self._tp.clear()
@@ -66,9 +74,9 @@ class ClassifyFPreRecMetric(Metric):
 
     def get_metric(self) -> dict:
         r"""
-        get_metric 函数将根据 update 函数累计的评价指标统计量来计算最终的评价结果.
+        :meth:`get_metric` 函数将根据 :meth:`update` 函数累计的评价指标统计量来计算最终的评价结果。
 
-        :return evaluate_result: {"acc": float}
+        :return: 包含以下内容的字典：``{"acc": float}``
         """
         evaluate_result = {}
 
@@ -127,13 +135,13 @@ class ClassifyFPreRecMetric(Metric):
 
     def update(self, pred, target, seq_len=None):
         r"""
-        update 函数将针对一个批次的预测结果做评价指标的累计
+        :meth:`update` 函数将针对一个批次的预测结果做评价指标的累计。
 
-        :param pred: 预测的 tensor, tensor 的形状可以是 [B,], [B, n_classes])
-                [B, max_len], 或者 [B, max_len, n_classes]
-        :param target: 真实值的 tensor, tensor 的形状可以是 [B,],
-                [B,], [B, max_len], 或者 [B, max_len]
-        :param seq_len: 序列长度标记, 标记的形状可以是 None, [B].
+        :param pred: 预测的 tensor, tensor 的形状可以是 ``[B,]`` 、``[B, n_classes]`` 、
+                ``[B, max_len]`` 或 ``[B, max_len, n_classes]``
+        :param target: 真实值的 tensor, tensor 的形状可以是 ``[B,]`` 、``[B,]`` 、``[B, max_len]``
+            或 ``[B, max_len]``
+        :param seq_len: 序列长度标记, 标记的形状可以是 ``None``,  或者 ``[B]``
 
         """
         pred = self.tensor2numpy(pred)

fastNLP/core/metrics/element.py

@@ -23,27 +23,27 @@ def _wrap_cal_value(func):
 
 
 class Element:
+    """
+    保存 :class:`~fastNLP.core.metrics.Metric` 中计算的元素值的对象
+
+    :param name: 名称
+    :param value: 元素的值
+    :param aggregate_method: 聚合的方法， 目前支持 ``['sum', 'mean', 'max', 'min']``:
+
+        * method 为 ``'sum'`` 时， 会将多张卡上聚合结果在维度为 `0` 上 累加起来。
+        * method 为 ``'mean'`` 时，会将多张卡上聚合结果在维度为 `0` 上取平均值。
+        * method 为 ``'max'`` 时，会将多张卡上聚合结果在维度为 `0` 上取最大值。
+        * method 为 ``'min'`` 时，会将多张卡上聚合结果在维度为 `0` 上取最小值。
+
+    :param backend: 使用的 backend 。Element 的类型会根据 ``backend`` 进行实际的初始化。例如 ``backend`` 为 ``'torch'`` 则该对象为
+        :class:`torch.Tensor` ； 如果 ``'backend'`` 为 ``'paddle'`` 则该对象为 :class:`paddle.Tensor` ；如果 ``backend`` 为
+        ``'jittor'`` , 则该对象为 :class:`jittor.Var` 。一般情况下直接默认为 ``'auto'`` 就行了， **fastNLP** 会根据实际调用 :meth`Metric.update`
+        函数时传入的参数进行合理的初始化，例如当传入的参数中只包含 :class:`torch.Tensor` 这一种 tensor 时（可以有其它非 tensor 类型的输入）
+        则认为 ``backend`` 为 ``'torch'`` ；只包含 :class:`jittor.Var` 这一种 tensor 时（可以有其它非 tensor 类型的输入）则认为 ``backend``
+        为 ``'jittor'`` 。如果没有检测到任何一种 tensor ，就默认使用 :class:`float` 类型作为 element 。
+
+    """
     def __init__(self, name, value: float, aggregate_method, backend: Backend):
-        """
-        保存 Metric 中计算的元素值的对象
-
-        :param name: 名称
-        :param value: 元素的值
-        :param aggregate_method: 聚合的方法， 目前支持 ``['sum', 'mean', 'max', 'mix']``:
-
-            * method 为 ``'sum'`` 时， 会将多张卡上聚合结果在维度为 `0` 上 累加起来。
-            * method 为 ``'mean'`` 时，会将多张卡上聚合结果在维度为 `0` 上取平均值。
-            * method 为 ``'max'`` 时，会将多张卡上聚合结果在维度为 `0` 上取最大值。
-            * method 为 ``'mix'`` 时，会将多张卡上聚合结果在维度为 `0` 上取最小值。
-
-        :param backend: 使用的 backend 。Element 的类型会根据 backend 进行实际的初始化。例如 backend 为 torch 则该对象为
-            Torch.tensor ； 如果backend 为 paddle 则该对象为 paddle.tensor ；如果 backend 为 jittor , 则该对象为 jittor.Var 。
-            一般情况下直接默认为 auto 就行了，fastNLP 会根据实际调用 Metric.update() 函数时传入的参数进行合理的初始化，例如当传入
-            的参数中只包含 torch.Tensor 这一种 tensor 时（可以有其它非 tensor 类型的输入）则认为 backend 为 torch ；只包含
-            jittor.Var 则认为 backend 这一种 tensor 时（可以有其它非 tensor 类型的输入）则认为 backend 为 jittor 。如果没有检测
-            到任何一种 tensor ，就默认使用 float 类型作为 element 。
-
-        """
         self.name = name
         self.init_value = value
         self.aggregate_method = aggregate_method
@@ -64,7 +64,6 @@ class Element:
     def aggregate(self):
         """
         自动 aggregate 对应的元素
-
         """
         self._check_value_initialized()
         if self.aggregate_method is None:  # 如果没有 aggregate 则不进行聚合。
@@ -116,7 +115,7 @@ class Element:
 
     def fill_value(self, value):
         """
-        对元素进行 fill_value， 会执行队友 backend 的 fill_value 方法
+        对元素进行 :meth:`fill_value` ， 会执行对应 backend 的 :meth:`fill_value` 方法
 
         """
         self._check_value_initialized()
@@ -136,7 +135,6 @@ class Element:
     def _check_value_initialized(self):
         """
         检查 Element 的 value 是否初始化了
-
         """
         if self._value is None:
             assert self.backend.is_specified(), f"Backend is not specified, please specify backend in the Metric " \
@@ -302,7 +300,7 @@ class Element:
 
     def __getattr__(self, item):
         """
-        为FDataLoader提供dataset的方法和属性，实现该方法后，用户可以在FDataLoader实例化后使用apply等dataset的方法
+        为 FDataLoader 提供 dataset 的方法和属性，实现该方法后，用户可以在 FDataLoader 实例化后使用 apply 等 dataset 的方法
         :param item:
         :return:
         """

fastNLP/core/metrics/metric.py

@@ -15,13 +15,13 @@ from fastNLP.core.metrics.element import Element
 
 class Metric:
     """
-    fastNLP 中 Metric 的基类，自定义 Metric 时，请继承该对象。使用该对象，将有助于减少在分布式状态下的 Metric 计算。
+    **fastNLP** 中 :class:`Metric` 的基类，自定义 :class:`Metric` 时，请继承该对象。使用该对象，将有助于减少在分布式状态下的 Metric 计算。
 
-    :param backend: 目前支持四种类型的 backend, ``[torch, paddle, jittor, auto]``。其中 ``auto`` 表示根据实际调用
-        Metric.update() 函数时传入的参数决定具体的 ``backend`` ，大部分情况下直接使用 ``auto`` 即可。
-    :param aggregate_when_get_metric: 在计算 metric 的时候是否自动将各个进程上的相同的 element 的数字聚合后再得到metric，
-        当 backend 不支持分布式时，该参数无意义。如果为 None ，将在 :class:`~fastNLP.core.controllers.Evaluator` 中根据 sampler 是否使用分布式
-        进行自动设置。
+    :param backend: 目前支持五种类型的 backend, ``['torch', 'paddle', 'jittor', 'oneflow', 'auto']``。其中 ``'auto'`` 表示根据实际调用 :meth:`update`
+        函数时传入的参数决定具体的 backend ，大部分情况下直接使用 ``'auto'`` 即可。
+    :param aggregate_when_get_metric: 在计算 metric 的时候是否自动将各个进程上的相同的 element 的数字聚合后再得到 metric，
+        当 backend 不支持分布式时，该参数无意义。如果为 ``None`` ，将在 :class:`~fastNLP.core.controllers.Evaluator` 中根据
+        sampler 是否使用分布式进行自动设置。
     """
     def __init__(self, backend: Union[str, Backend, None] = 'auto', aggregate_when_get_metric: bool = None):
         self.backend = AutoBackend(backend)
@@ -39,22 +39,22 @@ class Metric:
 
     def register_element(self, name, value: float = 0, aggregate_method=None, backend='auto') -> Element:
         """
-        注册一个 element 对象，注册之后便可以通过在 Metric 中直接通过 self.{name} 进行调用，可以认为该对象即为对应 backend 的
+        注册一个 element 对象，注册之后便可以通过在 Metric 中直接通过 ``self.{name}`` 进行调用，可以认为该对象即为对应 backend 的
         tensor 直接进行加减乘除计算即可。
 
-        ..warning::
+        .. warning::
 
-            如果想使得该 metric 可自动扩展到多卡的情况，请一定申明 aggregate_method 。
+            如果想使得该 metric 可自动扩展到多卡的情况，请一定申明 ``aggregate_method`` 。
 
-        :param name: 当前 element 的名字，注册后，在 Metric 中可以通过 self.{name} 访问该变量。
-        :param value: 初始化的值。在调用 Metric.reset() 方法时也将自动设置为该值
+        :param name: 当前 element 的名字，注册后，在 Metric 中可以通过 ``self.{name}`` 访问该变量。
+        :param value: 初始化的值。在调用 :meth:`Metric.reset` 方法时也将自动设置为该值
         :param aggregate_method: 如何聚合多卡上的结果，如果为单卡执行，该值无意义。如果设置为 None 则表示该 element 不进行聚合。
-        :param backend: 使用的 backend 。Element 的类型会根据 backend 进行实际的初始化。例如 backend 为 torch 则该对象为
-            Torch.tensor ； 如果backend 为 paddle 则该对象为 paddle.tensor ；如果 backend 为 jittor , 则该对象为 jittor.Var 。
-            一般情况下直接默认为 auto 就行了，fastNLP 会根据实际调用 Metric.update() 函数时传入的参数进行合理的初始化，例如当传入
-            的参数中只包含 torch.Tensor 这一种 tensor 时（可以有其它非 tensor 类型的输入）则认为 backend 为 torch ；只包含
-            jittor.Var 则认为 backend 这一种 tensor 时（可以有其它非 tensor 类型的输入）则认为 backend 为 jittor 。如果没有检测
-            到任何一种 tensor ，就默认使用 float 类型作为 element 。
+        :param backend: 使用的 backend 。Element 的类型会根据 ``backend`` 进行实际的初始化。例如 ``backend`` 为 ``'torch'`` 则该对象为
+            :class:`torch.Tensor` ； 如果 ``'backend'`` 为 ``'paddle'`` 则该对象为 :class:`paddle.Tensor` ；如果 ``backend`` 为
+            ``'jittor'`` , 则该对象为 :class:`jittor.Var` 。一般情况下直接默认为 ``'auto'`` 就行了， **fastNLP** 会根据实际调用 :meth`Metric.update`
+            函数时传入的参数进行合理的初始化，例如当传入的参数中只包含 :class:`torch.Tensor` 这一种 tensor 时（可以有其它非 tensor 类型的输入）
+            则认为 ``backend`` 为 ``'torch'`` ；只包含 :class:`jittor.Var` 这一种 tensor 时（可以有其它非 tensor 类型的输入）则认为 ``backend``
+            为 ``'jittor'`` 。如果没有检测到任何一种 tensor ，就默认使用 :class:`float` 类型作为 element 。
         :return: 注册的 Element 对象
         """
         if backend == 'auto':
@@ -71,8 +71,8 @@ class Metric:
 
     def reset(self):
         """
-        如果有非 element 的对象需要 reset 的时候，在本方法中写下非 element 的reset 方式。注册的 element 对象会自动 reset 为初始值。
-
+        在对每个 ``evaluate_dataloaders`` 遍历进行验证之前，:meth:`reset` 函数会被调用来重置每个非 element 对象；
+        如果有非 element 的对象需要重置的时候，在本方法中写下非 element 的重置方式。注册的 element 对象则会自动 reset 为初始值。
         """
         pass
 
@@ -142,9 +142,8 @@ class Metric:
     @contextmanager
     def sync(self, recover=True, aggregate=False):
         """
-        在这个上下文下， metric 会自动先同步需要同步操作的 element 。当 recover 为 True 时，在退出环境的时候，会重新将 element 的
-            值恢复到计算前的值。
-
+        在这个上下文下， :meth:`Metric` 会自动先同步需要同步操作的 element 。当 ``recover`` 为 ``True`` 时，在退出环境的时候，会重新将 element 的
+        值恢复到计算前的值。
         """
         keep_value = {}
         if aggregate:
@@ -172,14 +171,14 @@ class Metric:
 
     def set_auto_aggregate_when_get_metric(self, flag: bool):
         """
-        设置是否在 get_metric 的时候自动 aggregate
+        设置是否在 :meth:`get_metric` 的时候自动 aggregate
 
         """
         self.aggregate_when_get_metric = flag
 
     def tensor2numpy(self, tensor) -> np.array:
         """
-        将tensor向量转为numpy类型变量
+        将 ``tensor`` 向量转为 :class:`numpy.array` 类型变量。
 
         :param tensor:
         :return:
@@ -188,7 +187,7 @@ class Metric:
 
     def to(self, device):
         """
-        将所有的 element 变量移动到 device 设备上
+        将所有的 element 变量移动到 ``device`` 设备上
 
         :param device:
         :return:
@@ -198,7 +197,7 @@ class Metric:
 
     def all_gather_object(self, obj, group=None)->List:
         """
-        给定 obj 将各个 rank 上的 obj 汇总到每个 obj 上。返回一个 list 对象，里面依次为各个 rank 对应的 obj 。
+        给定 ``obj`` 将各个 rank 上的 ``obj`` 汇总到每个 ``obj`` 上。返回一个 list 对象，里面依次为各个 rank 对应的 ``obj`` 。
 
         :param obj: 需要汇总的对象，必须是个 pickable 的对象。
         :param group:

fastNLP/core/metrics/span_f1_pre_rec_metric.py

@@ -201,19 +201,31 @@ def _bio_tag_to_spans(tags, ignore_labels=None):
 
 class SpanFPreRecMetric(Metric):
     r"""
-
-    :param tag_vocab: 标签的 :class:`~fastNLP.Vocabulary` 。支持的标签为"B"(没有label)；或"B-xxx"(xxx为某种label，比如POS中的NN)，
-        在解码时，会将相同xxx的认为是同一个label，比如['B-NN', 'E-NN']会被合并为一个'NN'.
-    :param encoding_type: 目前支持bio, bmes, bmeso, bioes。默认为None，通过tag_vocab自动判断.
-    :param ignore_labels: str 组成的list. 这个list中的class不会被用于计算。例如在POS tagging时传入['NN']，则不会计算'NN'个label
-    :param only_gross: 是否只计算总的f1, precision, recall的值；如果为False，不仅返回总的f1, pre, rec, 还会返回每个label的f1, pre, rec
-    :param f_type: `micro` 或 `macro` . `micro` :通过先计算总体的TP，FN和FP的数量，再计算f, precision, recall; `macro` : 分布计算每个类别的f, precision, recall，然后做平均（各类别f的权重相同）
-    :param beta: f_beta分数， :math:`f_{beta} = \frac{(1 + {beta}^{2})*(pre*rec)}{({beta}^{2}*pre + rec)}` . 常用为 `beta=0.5, 1, 2` 若为0.5则精确率的权重高于召回率；若为1，则两者平等；若为2，则召回率权重高于精确率。
-    :param backend: 目前支持四种类型的 backend, ``[torch, paddle, jittor, auto]``。其中 ``auto`` 表示根据实际调用
-        Metric.update() 函数时传入的参数决定具体的 ``backend`` ，大部分情况下直接使用 ``auto`` 即可。
-    :param aggregate_when_get_metric: 在计算 metric 的时候是否自动将各个进程上的相同的 element 的数字聚合后再得到metric，
-        当 backend 不支持分布式时，该参数无意义。如果为 None ，将在 :class:`~fastNLP.core.controllers.Evaluator` 中根据 sampler 是否使用分布式
-        进行自动设置。
+    在 **序列标注** 任务中评估抽取结果匹配度的 **Metric** 。
+
+    :param tag_vocab: 标签的 :class:`~fastNLP.core.Vocabulary` 。支持的标签有 ``"B"`` (没有label)；或 ``"B-xxx"`` ( ``xxx`` 为某种 label ，比如 POS 中的 NN)，
+        在解码时，会将相同 ``xxx`` 的认为是同一个 label ，比如 ['B-NN', 'E-NN'] 会被合并为一个 'NN' 。
+    :param encoding_type: 目前支持 ``['bio', 'bmes', 'bmeso', 'bioes', None]`` 。默认为 ``None`` ，通过 ``tag_vocab`` 自动判断
+    :param ignore_labels: 字符串组成的列表，这个列表中包含的内容不会被用于计算。例如在 *POS tagging* 时传入 ``['NN']`` ，则不会计算 ``'NN'`` 这个 ``label``
+    :param only_gross: 是否只计算总的 ``f1``, ``precision`` , ``recall`` 的值。如果为 ``False`` ，不仅返回总的 ``f1`` , ``pre`` , ``rec`` , 还会返回每个 label 
+        ``f1`` , ``pre`` , ``rec`` 
+    :param f_type: ``'micro'`` 或 ``'macro'``。
+    
+        - *micro* -- 通过先计算总体的 ``TP``， ``FN`` 和 ``FP`` 的数量，再计算 ``f``, ``precision``, ``recall`` ；
+        - *macro* -- 分别计算每个类别的 ``f`` , ``precision`` , ``recall`` ，然后做平均（各类别 ``f`` 的权重相同）；
+    
+    :param beta: **f_beta** 分数中的 ``beta`` 值。 常用为 ``beta=0.5, 1, 2`` 若为 0.5 则 **精确率** 的权重高于 **召回率** ；若为1，则两者平等；若为2，则
+        **召回率** 权重高于 **精确率** 。**f_beta** 分数的计算公式为：
+        
+        .. math::
+
+            f_{beta} = \\frac{(1 + {beta}^{2})*(pre*rec)}{({beta}^{2}*pre + rec)}
+
+    :param backend: 目前支持五种类型的 backend, ``['torch', 'paddle', 'jittor', 'oneflow', 'auto']``。其中 ``'auto'`` 表示根据实际调用 :meth:`update`
+        函数时传入的参数决定具体的 backend ，大部分情况下直接使用 ``'auto'`` 即可。
+    :param aggregate_when_get_metric: 在计算 metric 的时候是否自动将各个进程上的相同的 element 的数字聚合后再得到 metric，
+        当 backend 不支持分布式时，该参数无意义。如果为 ``None`` ，将在 :class:`~fastNLP.core.controllers.Evaluator` 中根据
+        sampler 是否使用分布式进行自动设置。
     """
     def __init__(self, tag_vocab: Vocabulary, encoding_type: str = None, ignore_labels: List[str] = None,
                  only_gross: bool = True, f_type='micro',
@@ -262,7 +274,7 @@ class SpanFPreRecMetric(Metric):
 
     def get_metric(self) -> dict:
         """
-        get_metric 函数将根据 update 函数累计的评价指标统计量来计算最终的评价结果.
+        :meth:`get_metric` 函数将根据 :meth:`update` 函数累计的评价指标统计量来计算最终的评价结果。
         """
         evaluate_result = {}
 
@@ -317,12 +329,12 @@ class SpanFPreRecMetric(Metric):
         return evaluate_result
 
     def update(self, pred, target, seq_len: Optional[List] = None) -> None:
-        r"""u
-        pdate函数将针对一个批次的预测结果做评价指标的累计
+        r"""
+        :meth:`update` 函数将针对一个批次的预测结果做评价指标的累计。
 
-        :param pred: [batch, seq_len] 或者 [batch, seq_len, len(tag_vocab)], 预测的结果
-        :param target: [batch, seq_len], 真实值
-        :param seq_len: [batch] 文本长度标记
+        :param pred: 预测的结果，大小为 ``[batch, seq_len]`` 或者 ``[batch, seq_len, len(tag_vocab)]``
+        :param target: 真实值，大小为 ``[batch, seq_len]``
+        :param seq_len: 文本长度标记，大小为 ``[batch]``
         :return:
         """
         pred = self.tensor2numpy(pred)

fastNLP/core/samplers/conversion_utils.py

@@ -7,10 +7,10 @@ from fastNLP.core.samplers.unrepeated_sampler import UnrepeatedSampler, Unrepeat
 
 def conversion_between_reproducible_and_unrepeated_sampler(sampler):
     """
-    将 sampler 替换成其对应的 reproducible 版本或 unrepeated 版本。如果输入是 UnrepeatedSampler 但是没找到对应的
-    ReproducibleSampler，
+    将 ``sampler`` 替换成其对应的 reproducible 版本或 unrepeated 版本。如果输入是 :class:`~fastNLP.core.samplers.unrepeated_sampler.UnrepeatedSampler`
+    但是没找到对应的 :class:`~fastNLP.core.samplers.reproducible_sampler.ReproducibleSampler` 则会报错。
 
-    :param sampler: 需要转换的 sampler 。
+    :param sampler: 需要转换的 ``sampler`` ；
     :return:
     """
     assert isinstance(sampler, UnrepeatedSampler) or isinstance(sampler, ReproducibleSampler), \

fastNLP/core/samplers/mix_sampler.py

@@ -18,20 +18,30 @@ if _NEED_IMPORT_TORCH:
 
 class MixSampler:
     """
-    mix_sampler的基类
+    所有 mix_sampler 的基类。
+
+    :param dataset: 一个字典，每个元素都是一个实现了 __getitem__ 和 __len__ 的数据容器
+    :param batch_size: ``dataset`` 的批次大小，所有 ``dataset`` 均采用该 ``batch_size`` 作为批次大小
+    :param sampler: 实例化好的 ``sampler`` ，每个 ``dataset`` 对应一个 ``sampler`` 对象
+    :param ds_ratio: ``ds_ratio`` 是控制 datasets 怎么组成一个混合大数据集的重要参数， 其取值为 ``[None, 'truncate_to_least', 'pad_to_most', List[float], Dict[str, float]]``:
+
+        * ds_ratio 为 ``None``, datasets 数据集序列或字典不进行数据扩充处理；
+        * ds_ratio 为 ``'truncate_to_least'``, datasets 数据集序列或字典会计算得到 datasets序列中 dataset 最断长度 ``mix_len``， 其他数据集会被切断
+          到最短长度 ``mix_len``。这种切断不是物理上切断，``MixDataLoader`` 会根据 sampler 不同来采样数据集到指定的最短长度 ``mix_len``；
+        * ds_ratio 为 ``'pad_to_most'``, datasets 数据集序列或字典会计算得到 datasets序列中 dataset 最大长度 ``max_len``, 其他其他数据集会扩充
+          到最大长度 ``mix_len``。这种扩充不是物理上扩充， ``MixDataLoader`` 会根据 sampler 不同来重采样 dataset 到指定的最大长度 ``max_len``；
+        * ds_ratio 为 ``Dict[str, float]`` 时， datasets 类型也必须为 ``Dict[str, DataSet]``, 其 key 一一对应。 ``ds_ratio`` 的 value 是任意大于 0 的浮点数，
+          代表着 datasets 的 value 数据进行扩充或者缩减的倍数；
+
+    :param drop_last: 当最后一个 batch 长度小于 ``batch_size`` 时是否丢弃
+    :param rank: 分布式训练中当前进程的 ``global_rank``
+    :param world_size: 分布式训练中进程的总数 **world_size**
     """
 
     def __init__(self, dataset: Dict, batch_size: int = None,
                  sampler: Union[Dict[str, "Sampler"], None, str] = None,
                  ds_ratio: Union[str, Dict[str, float]] = None,
                  drop_last: bool = False, rank: int = -1, word_size: int = -1) -> None:
-        """
-
-        :param dataset: 实现了__getitem__和__len__的数据容器列表
-        :param batch_size: 对应dataset的批次大小，可以为list或者为int，当为int时默认所有dataset
-        :param sampler: 实例化好的sampler，每个dataset对应一个sampler对象
-        :param drop_last: 是否去掉最后一个batch的数据，其长度小于batch_size
-        """
         # sampler 为 dict，则判断是否与 datasets 的 key 相同
         if isinstance(sampler, Dict):
             for key in dataset.keys():
@@ -119,7 +129,7 @@ class MixSampler:
 
 class InnerSampler:
     """
-    提供多卡情况下使用的内部sampler
+    提供多卡情况下使用的内部 sampler
     """
     def __init__(self, ds_ind_list: List) -> None:
         self.ds_ind_list = ds_ind_list
@@ -134,7 +144,8 @@ class InnerSampler:
 
 class DopedSampler(MixSampler):
     """
-    定制给MixDataLoader的BatchSampler，其功能是将传入的datasets的list列表混合采样组成一个个batch返回。
+    定制给 :class:`~fastNLP.core.dataloaders.MixDataLoader` 的 ``BatchSampler``，其功能是将传入的 ``datasets`` 
+    字典混合采样组成一个个 batch 返回。
     """
     def __init__(self, dataset: Dict, batch_size: int = None,
                  sampler: Union[Dict[str, "Sampler"], str] = None,
@@ -256,27 +267,21 @@ class DopedSampler(MixSampler):
 
 class MixSequentialSampler(MixSampler):
     """
-    定制给MixDataLoader的BatchSampler，其功能是将传入的datasets的list列表顺序采样并返回index，只有处理了上一个dataset才会处理下一个。
+    定制给 :class:`~fastNLP.core.dataloaders.MixDataLoader` 的 ``BatchSampler``，其功能是将传入的 ``datasets`` 按顺序采样并返回 index，
+    只有上一个 dataset 处理结束后才会处理下一个。
     """
 
-    def __init__(self, dataset: Union[List, Dict], batch_size: int = None,
+    def __init__(self, dataset: Dict, batch_size: int = None,
                  sampler: Union[List["Sampler"], Dict[str, "Sampler"], None, str] = None,
                  ds_ratio: Union[str, List[float], Dict[str, float]] = None,
                  drop_last: bool = False, rank: int = -1, word_size: int = -1) -> None:
-        """
-
-        :param dataset: 实现了__getitem__和__len__的数据容器列表
-        :param batch_size: 对应dataset的批次大小，可以为list或者为int，当为int时默认所有dataset
-        :param sampler: 实例化好的sampler，每个dataset对应一个sampler对象
-        :param drop_last: 是否去掉最后一个batch的数据，其长度小于batch_size
-        """
         super(MixSequentialSampler, self).__init__(dataset=dataset, batch_size=batch_size,
                                                    sampler=sampler, ds_ratio=ds_ratio,
                                                    drop_last=drop_last, rank=rank, word_size=word_size)
 
     def __iter__(self) -> Iterable[List[int]]:
         """
-        按照dataset的顺序采样，打包成一个batch后返回
+        按照 ``dataset`` 的顺序采样，打包成一个 batch 后返回。
 
         :return:
         """
@@ -384,22 +389,14 @@ class MixSequentialSampler(MixSampler):
 
 class PollingSampler(MixSampler):
     """
-    定制给MixDataLoader的BatchSampler，其功能是将传入的datasets的list列表轮流采样并返回index，处理了上个dataset的一个batch后会处理下一个。
+    定制给 :class:`~fastNLP.core.dataloaders.MixDataLoader` 的 ``BatchSampler``，其功能是将传入的 ``datasets`` 轮流采样并返回 index，
+    处理结束上个 dataset 的一个 batch 后会处理下一个。
     """
 
     def __init__(self, dataset: Union[List, Dict], batch_size: int = 16,
                  sampler: Union[List["Sampler"], Dict[str, "Sampler"], str] = None,
                  drop_last: bool = False, ds_ratio="pad_to_most", rank: int = -1,
                  word_size: int = -1) -> None:
-        """
-
-        :param dataset: 实现了__getitem__和__len__的数据容器列表
-        :param batch_size: 对应dataset的批次大小，可以为list或者为int，当为int时默认所有dataset
-        :param sampler: 实例化好的sampler，每个dataset对应一个sampler对象
-        :param drop_last: 是否去掉最后一个batch的数据，其长度小于batch_size
-        :param ds_ratio: 当ds_ratio=None时候， 轮流采样dataset列表直至所有的数据集采样完；当ds_ratio='truncate_to_least'时，
-            以dataset列表最短的ds为基准，长的数据集会被截断；当ds_ratio='pad_to_most'时，以dataset列表最长ds为基准，短的数据集会被重采样
-       """
         super(PollingSampler, self).__init__(dataset=dataset, batch_size=batch_size,
                                              sampler=sampler, ds_ratio=ds_ratio,
                                              drop_last=drop_last, rank=rank, word_size=word_size)

fastNLP/core/samplers/reproducible_batch_sampler.py

@@ -1,3 +1,14 @@
+"""
+:class:`ReproducibleBatchSampler` 是 **fastNLP** 提供的一种特殊 BatchSampler，它可以记录采样过程中每一次采样和 epoch 的信息，
+方便在保存-加载后能够从上一次采样结束的地方继续进行采样，实现 **断点重训** 。
+
+.. note::
+
+    DataLoader 中只要存在 :class:`~fastNLP.core.samplers.reproducible_sampler.ReproducibleSampler` 或 :class:`ReproducibleBatchSampler`
+    中的一个便可以实现断点重训复现的功能。
+
+"""
+
 __all__ = [
     'BucketedBatchSampler',
     "ReproduceBatchSampler",
@@ -19,6 +30,12 @@ from abc import abstractmethod
 
 
 class ReproducibleBatchSampler:
+    """
+    **可复现**的 BatchSampler 对象。
+
+    注意所有继承 :class:`ReproducibleBatchSampler` 的类的 :meth:`__init__` 方法中都需要加入参数 `**kwargs`，用来使我们再断点重训时重新实例化这个 BatchSampler
+    注意，所有 :meth:`__init__` 中初始化的变量，都不能含有 ``_`` 下横线作为开头；所有不在 :meth:`__init__` 中设置的变量都必须以下横线开头。
+    """
     def __init__(self, **kwargs):
         self.num_replicas = 1
 
@@ -57,13 +74,13 @@ class ReproducibleBatchSampler:
 
 class ReproduceBatchSampler(ReproducibleBatchSampler):
     """
-    可以使得 batch_sampler 对象状态恢复的 wrapper 。
+    可以使得 ``batch_sampler`` 对象状态恢复的 wrapper 。
 
-    :param batch_sampler: 可迭代出 数字 或 数字列表 的可迭代对象。ReproduceBatchSampler 将首先遍历一边该对象，然后将迭代
-        出来的序号暂存起来，使用时按照 batch_size 的 batch 大小吐出序号列表。
-    :param batch_size: 每个 batch 的大小是多少。
-    :param drop_last: 如果最后一个 batch 无法构成 batch_size 那么多个 sample ，是否丢掉。
-    :param kwargs: fastNLP 内部使用。
+    :param batch_sampler: 可迭代出 **数字** 或 **数字列表** 的可迭代对象。:class:`ReproduceBatchSampler` 将首先遍历一边该对象，然后将迭代
+        出来的序号暂存起来，使用时按照 ``batch_size`` 的 batch 大小吐出序号列表。
+    :param batch_size: 每个 batch 的大小是多少
+    :param drop_last: 如果最后一个 batch 无法构成 ``batch_size`` 个 sample ，是否丢掉
+    :param kwargs: fastNLP 内部使用的参数
     """
     def __init__(self, batch_sampler, batch_size: int, drop_last: bool, **kwargs):
         super().__init__()
@@ -162,12 +179,12 @@ class RandomBatchSampler(ReproducibleBatchSampler):
     """
     随机分 batch 的 batch_sampler 。
 
-    :param dataset: 实现了 __len__ 方法的数据容器。
+    :param dataset: 实现了 __len__ 方法的数据容器
     :param batch_size: 每个 batch 的大小
-    :param shuffle: 如果为 True，将不进行 shuffle，实际上数据会以从长到短的方式输出。
-    :param drop_last: 如果最后一个 batch 的 sample 数量无法凑齐 batch_size 这么多，是否需要丢掉。
+    :param shuffle: 如果为 ``True``，将不进行打乱操作，实际上数据会以从长到短的方式输出
+    :param drop_last: 如果最后一个 batch 无法构成 batch_size 个 sample ，是否丢掉
     :param seed: 设置的随机数种子
-    :param kwargs: fastNLP 保留使用
+    :param kwargs: fastNLP 内部使用的参数
     """
     def __init__(self, dataset, batch_size:int = 32, shuffle: bool = True,
                  drop_last: bool = False, seed: int = 0, **kwargs):
@@ -195,6 +212,15 @@ class RandomBatchSampler(ReproducibleBatchSampler):
         self.old_batch_size = kwargs.get('old_batch_size', self.batch_size)
 
     def set_distributed(self, num_replicas, rank, pad=True):
+        """
+        进行分布式的相关设置，应当在初始化该 BatchSampler 本身后立即被调用。
+
+        :param num_replicas: 分布式训练中的进程总数
+        :param rank: 当前进程的 ``global_rank``
+        :param pad: 如果 sample 数量不整除 ``num_replicas`` 的时候，要不要 pad 一下，使得最终使得每个进程上
+            的 sample 数量是完全一致的
+        :return: 自身
+        """
         assert self.during_iter is False, "Cannot set the sampler to be distributed when it is " \
                                           "during an unfinished iteration."
         assert num_replicas > 0 and isinstance(num_replicas, int)
@@ -266,14 +292,14 @@ class RandomBatchSampler(ReproducibleBatchSampler):
         if self.epoch < 0:  # 防止用户没有修改epoch，导致每个epoch都一样了
             self.epoch -= 1
 
-    def batchify(self, indices, batch_size, seed):
+    def batchify(self, indices, batch_size, seed) -> List[List[int]]:
         """
-        将 indices 分为 batches
+        将 ``indices`` 分为 batches
 
-        :param sorted_indices: List[int]
+        :param indices: List[int]
         :param batch_size: int
         :param seed: int
-        :return:  List[List[int]]
+        :return:
         """
         # 实际的 bucket 大小
         rng = np.random.default_rng(abs(seed))
@@ -299,19 +325,15 @@ class RandomBatchSampler(ReproducibleBatchSampler):
     @property
     def total_size(self):
         """
-        这个变量代表的含义是当前这个sampler会最终产生出的index数量（包括了其它rank的），因为replica和pad的原因，这个值可能等于、
-        大于或者小于len(dataset)
-
-        :return:
+        当前 BatchSampler 会最终产生出的 index 数量（包括了其它 rank 的），因为 ``replica`` 和 ``pad`` 的原因，这个值可能等于、
+        大于或者小于 ``len(dataset)``。
         """
         return self.num_consumed_samples + self.num_replicas*self.num_left_samples
 
     @property
     def num_left_samples(self):
         """
-        返回当前 iteration 还有多少个 sample 结束，表示的是当前 rank 的还剩多少。
-
-        :return:
+        当前迭代还有多少个 sample 结束，表示的是 **当前 rank** 的还剩多少。
         """
         num_consumed_samples = self.num_consumed_samples
         return math.ceil((self.num_samples - num_consumed_samples) / self.num_replicas) if \
@@ -320,9 +342,7 @@ class RandomBatchSampler(ReproducibleBatchSampler):
     @property
     def num_samples(self):
         """
-        返回样本的总数
-
-        :return:
+        样本的总数
         """
         total_len = getattr(self.dataset, 'total_len', None)
         if not isinstance(total_len, int):
@@ -377,18 +397,19 @@ class RandomBatchSampler(ReproducibleBatchSampler):
 
 class BucketedBatchSampler(ReproducibleBatchSampler):
     """
-    首先按照 ``sample`` 的长度排序，然后按照 batch_size*num_batch_per_bucket 为一个桶的大小，``sample`` 只会在这个桶内进行组
+    首先按照 ``sample`` 的长度排序，然后按照 *batch_size*num_batch_per_bucket* 为一个桶的大小，``sample`` 只会在这个桶内进行组
     合，这样每个 ``batch`` 中的 ``padding`` 数量会比较少 （因为桶内的数据的长度都接近）。
 
     :param dataset: 实现了 __len__ 方法的数据容器。
     :param length: 每条数据的长度。
 
         * 为 ``List[int]`` 时
-         应当与 dataset 有一样的长度，表示 dataset 中每个元素的数量；
+          应当与 dataset 有一样的长度，表示 dataset 中每个元素的数量；
         * 为 ``str`` 时
-         仅当传入的 ``dataset`` 是 :class:`~fastNLP.DataSet` 时，允许传入 `str` ，该 `str` 将被认为是 ``dataset`` 中的
+          仅当传入的 ``dataset`` 是 :class:`~fastNLP.DataSet` 时，允许传入 `str` ，该 `str` 将被认为是 ``dataset`` 中的
           ``field`` 。若 field 中的元素为 ``int``，则认为该值是 sample 的长度；若不为 ``int`` ，则尝试使用 ``len`` 方法
           获取该 ``field`` 中每个元素的长度。
+
     :param batch_size: 每个 batch 的大小
     :param num_batch_per_bucket: 多少个 ``batch`` 组成一个桶，数据只会在一个桶内进行 ``shuffle`` 。
     :param shuffle: 如果为 True，将不进行 ``shuffle``，实际上数据会以从长到短的方式输出。
@@ -440,6 +461,15 @@ class BucketedBatchSampler(ReproducibleBatchSampler):
         self.old_num_batch_per_bucket = kwargs.get('old_num_batch_per_bucket', self.num_batch_per_bucket)
 
     def set_distributed(self, num_replicas, rank, pad=True):
+        """
+        进行分布式的相关设置，应当在初始化该 BatchSampler 本身后立即被调用。
+
+        :param num_replicas: 分布式训练中的进程总数
+        :param rank: 当前进程的 ``global_rank``
+        :param pad: 如果 sample 数量不整除 ``num_replicas`` 的时候，要不要 pad 一下，使得最终使得每个进程上
+            的 sample 数量是完全一致的
+        :return:
+        """
         assert self.during_iter is False, "Cannot set the sampler to be distributed when it is " \
                                            "during an unfinished iteration."
         assert num_replicas > 0 and isinstance(num_replicas, int)
@@ -462,19 +492,15 @@ class BucketedBatchSampler(ReproducibleBatchSampler):
     @property
     def total_size(self):
         """
-        这个变量代表的含义是当前这个sampler会最终产生出的index数量（包括了其它rank的），因为replica和pad的原因，这个值可能等于、
-        大于或者小于len(dataset)
-
-        :return:
+        当前 BatchSampler 会最终产生出的 index 数量（包括了其它 rank 的），因为 ``replica`` 和 ``pad`` 的原因，这个值可能等于、
+        大于或者小于 ``len(dataset)``。
         """
         return self.num_consumed_samples + self.num_replicas*self.num_left_samples
 
     @property
     def num_left_samples(self):
         """
-        返回当前 iteration 还有多少个 sample 结束，表示的是当前 rank 的还剩多少。
-
-        :return:
+        当前迭代还有多少个 sample 结束，表示的是 **当前 rank** 的还剩多少。
         """
         num_consumed_samples = self.num_consumed_samples
         return math.ceil((self.num_samples - num_consumed_samples) / self.num_replicas) if \
@@ -483,9 +509,7 @@ class BucketedBatchSampler(ReproducibleBatchSampler):
     @property
     def num_samples(self):
         """
-        返回样本的总数
-
-        :return:
+        样本的总数
         """
         total_len = getattr(self.dataset, 'total_len', None)
         if not isinstance(total_len, int):
@@ -572,15 +596,15 @@ class BucketedBatchSampler(ReproducibleBatchSampler):
         if self.epoch < 0:  # 防止用户没有修改epoch，导致每个epoch都一样了
             self.epoch -= 1
 
-    def bucketerize(self, sorted_indices, batch_size, num_batch_per_bucket, seed):
+    def bucketerize(self, sorted_indices, batch_size, num_batch_per_bucket, seed) -> List[List[int]]:
         """
-        将 indices 分桶
+        将 ``indices`` 分桶
 
         :param sorted_indices: List[int]
         :param batch_size: int
         :param num_batch_per_bucket: int
         :param seed: int
-        :return:  List[List[int]]
+        :return:
         """
         # 实际的 bucket 大小
         bucket_size = min(len(sorted_indices), batch_size * num_batch_per_bucket)

fastNLP/core/samplers/reproducible_sampler.py

@@ -1,3 +1,14 @@
+"""
+:class:`ReproducibleSampler` 是 **fastNLP** 提供的一种特殊 Sampler，它可以记录采样过程中每一次采样和 epoch 的信息，
+方便在保存-加载后能够从上一次采样结束的地方继续进行采样，实现 **断点重训** 。
+
+.. note::
+
+    DataLoader 中只要存在 :class:`ReproducibleSampler` 或 :class:`~fastNLP.core.samplers.reproducible_batch_sampler.ReproducibleBatchSampler`
+    中的一个便可以实现断点重训复现的功能。
+
+"""
+
 __all__ = [
     'ReproducibleSampler',
     'RandomSampler',
@@ -16,10 +27,10 @@ from fastNLP.core.dataset import DataSet
 
 class ReproducibleSampler:
     """
-    可复现的 Sampler 对象。
+    **可复现** 的 Sampler 对象。
 
-    注意所有继承 `ReproducibleSampler` 的类的  `__init__` 方法中都需要加入参数 `**kwargs`，用来使我们再断点重训时重新实例化这个 sampler
-    或者 batch_sampler；注意，所有在 init 中初始化的变量，都不能含有 _ 下横线作为开头；所有不在 init 中设置的变量都必须以下横线开头。
+    注意所有继承 :class:`ReproducibleSampler` 的类的 :meth:`__init__` 方法中都需要加入参数 `**kwargs`，用来使我们再断点重训时重新实例化这个 Sampler
+    注意，所有 :meth:`__init__` 中初始化的变量，都不能含有 ``_`` 下横线作为开头；所有不在 :meth:`__init__` 中设置的变量都必须以下横线开头。
 
     """
     def __init__(self, **kwargs):
@@ -61,9 +72,9 @@ class RandomSampler(ReproducibleSampler):
     随机顺序的 Sampler 。
 
     :param dataset: 实现了 __len__ 方法的数据容器
-    :param shuffle: 是否在每次 iterate 的时候打乱顺序。
-    :param seed: 随机数种子。
-    :param kwargs: 用户不需要使用，fastNLP 内部使用
+    :param shuffle: 是否在每次 iterate 的时候打乱顺序
+    :param seed: 随机数种子
+    :param kwargs: fastNLP 内部使用的参数
     """
     def __init__(self, dataset, shuffle: bool = True, seed: int = 0, **kwargs):
         super(RandomSampler, self).__init__()
@@ -84,15 +95,16 @@ class RandomSampler(ReproducibleSampler):
 
     def __len__(self):
         """
-        返回 sampler 一次完整的迭代过程会产生多少个index。多卡的情况下，只考虑当前rank；
-        :return:
+        返回 sampler 一次完整的迭代过程会产生多少个index。多卡的情况下，只考虑当前rank。
         """
         return self.total_size//self.num_replicas
 
     def __iter__(self):
         r"""
-        当前使用num_consumed_samples做法会在交替使用的时候遇到问题；
+        当前使用 num_consumed_samples 做法会在交替使用的时候遇到问题。
+
         Example::
+
             >>> sampler = RandomSampler()
             >>> iter1 = iter(sampler)
             >>> iter2 = iter(sampler)
@@ -131,8 +143,6 @@ class RandomSampler(ReproducibleSampler):
     def generate_indices(self) -> List[int]:
         """
         生成随机序列
-
-        :return:
         """
         if self.shuffle:
             indices = list(range(self.num_samples))
@@ -176,12 +186,13 @@ class RandomSampler(ReproducibleSampler):
 
     def set_distributed(self, num_replicas:int, rank:int, pad:bool=True):
         """
+        进行分布式的相关设置，应当在初始化该 Sampler 本身后立即被调用。
 
-        :param num_replicas:
-        :param rank:
-        :param pad: 这个 pad 的意思是指如果 sample 数量不整除 num_replicas 的时候，要不要 pad 一下，使得最终使得 replica 上
-            的 sample 数量是完全一致的。
-        :return:
+        :param num_replicas: 分布式训练中的进程总数
+        :param rank: 当前进程的 ``global_rank``。
+        :param pad: 如果 sample 数量不整除 ``num_replicas`` 的时候，要不要 pad 一下，使得最终使得每个进程上
+            的 sample 数量是完全一致的
+        :return: 自身
         """
 
         assert self.during_iter is False, "Cannot set the sampler to be distributed when it is " \
@@ -198,18 +209,15 @@ class RandomSampler(ReproducibleSampler):
     @property
     def total_size(self):
         """
-        这个变量代表的含义是当前这个sampler会最终产生出的index数量，因为replica和pad的原因，这个值可能等于、大于或者小于len(dataset)
-
-        :return:
+        当前 sampler 会最终产生出的 index 数量（包括了其它 rank 的），因为 ``replica`` 和 ``pad`` 的原因，这个值可能等于、
+        大于或者小于 ``len(dataset)``。
         """
         return self.num_consumed_samples + self.num_replicas*self.num_left_samples
 
     @property
     def num_left_samples(self):
         """
-        返回当前 iteration 还有多少个 sample 结束。表示的是当前 rank 的还剩多少
-
-        :return:
+        当前迭代还有多少个 sample 结束。表示的是 **当前 rank** 的还剩多少
         """
         num_consumed_samples = self.num_consumed_samples
         return math.ceil((self.num_samples - num_consumed_samples) / self.num_replicas) if \
@@ -218,9 +226,7 @@ class RandomSampler(ReproducibleSampler):
     @property
     def num_samples(self):
         """
-        返回样本的总数
-
-        :return:
+        样本的总数
         """
         total_len = getattr(self.dataset, 'total_len', None)
         if not isinstance(total_len, int):
@@ -229,7 +235,7 @@ class RandomSampler(ReproducibleSampler):
 
 class SequentialSampler(RandomSampler):
     """
-    按照顺序读取 ``dataset`` 。在多卡情况下，间隔读取，例如，在两卡情况下，卡 0 取 ``[0,2,4,..]``, 卡1取 ``[1,3,5...]`` 。
+    按照顺序读取 ``dataset`` 。在多卡情况下，间隔读取，例如，在两卡情况下，卡 0 取 ``[0,2,4,..]``, 卡 1 取 ``[1,3,5...]`` 。
 
     :param dataset: 实现了 __len__ 方法的数据容器。
     :param kwargs:
@@ -300,17 +306,18 @@ class SortedSampler(SequentialSampler):
     将 ``dataset`` 中的数据根据 ``length`` 从长到短进行迭代。在多卡情况下，由于 ``padding`` , 最后一个 ``sample`` 可能是最长
     的那个 ``sample`` 。
 
-    :param dataset: 实现了 __len__ 方法的数据容器。
-    :param length: 每条数据的长度。
+    :param dataset: 实现了 __len__ 方法的数据容器
+    :param length: 每条数据的长度：
 
         * 为 ``List[int]`` 时
-         应当与 dataset 有一样的长度，表示 dataset 中每个元素的数量；
+          应当与 dataset 有一样的长度，表示 dataset 中每个元素的数量；
         * 为 ``str`` 时
-         仅当传入的 ``dataset`` 是 :class:`~fastNLP.DataSet` 时，允许传入 `str` ，该 `str` 将被认为是 ``dataset`` 中的
+          仅当传入的 ``dataset`` 是 :class:`~fastNLP.DataSet` 时，允许传入 `str` ，该 `str` 将被认为是 ``dataset`` 中的
           ``field`` 。若 field 中的元素为 ``int``，则认为该值是 sample 的长度；若不为 ``int`` ，则尝试使用 ``len`` 方法
-          获取该 ``field`` 中每个元素的长度。
-    :param seed: 设置的随机数种子。
-    :param kwargs: fastNLP 保留使用。
+          获取该 ``field`` 中每个元素的长度；
+
+    :param seed: 设置的随机数种子
+    :param kwargs: fastNLP 内部使用的参数
     """
     def __init__(self, dataset, length:Union[str, List], **kwargs):
         super().__init__(dataset=dataset, **kwargs)

fastNLP/core/samplers/unrepeated_sampler.py

@@ -13,19 +13,19 @@ import numpy as np
 
 class UnrepeatedSampler:
     """
-    在多卡场景下保证 indice 不重复的 sampler
+    在多卡场景下保证 indice 不重复的 Sampler。
     """
     pass
 
 
 class UnrepeatedRandomSampler(UnrepeatedSampler):
     """
-    考虑在多卡 evaluate 的场景下，不能重复 sample。
+    考虑在多卡 evaluate 的场景下，不能重复采样。
 
-    :param dataset: 实现了 __len__ 方法的数据容器。
-    :param shuffle: 如果为 True，将不进行 shuffle，实际上数据会以从长到短的方式输出。
+    :param dataset: 实现了 __len__ 方法的数据容器
+    :param shuffle: 如果为 ``True``，将不进行 shuffle，实际上数据会以从长到短的方式输出
     :param seed: 设置的随机数种子
-    :param kwargs: fastNLP 保留使用
+    :param kwargs: fastNLP 内部使用的参数
     """
     def __init__(self, dataset, shuffle: bool = False, seed: int = 0, **kwargs):
         self.dataset = dataset
@@ -39,7 +39,7 @@ class UnrepeatedRandomSampler(UnrepeatedSampler):
 
     def __len__(self):
         """
-        返回 sampler 一次完整的迭代过程会产生多少个index。多卡的情况下，只考虑当前rank；
+        返回 ``Sampler`` 一次完整的迭代过程会产生多少个 index 。多卡的情况下，只考虑 **当前rank** 。
         :return:
         """
         num_common = self.num_samples//self.num_replicas
@@ -78,11 +78,11 @@ class UnrepeatedRandomSampler(UnrepeatedSampler):
 
     def set_distributed(self, num_replicas, rank):
         """
-        该方法本质上等同于 ddp 情形下的没有完成的初始化，应当在初始化该 sampler 本身后立即被调用；
+        该方法本质上等同于 ddp 情形下的没有完成的初始化，应当在初始化该 Sampler 本身后立即被调用。
 
-        :param num_replicas:
-        :param rank:
-        :return:
+        :param num_replicas: 分布式训练中的进程总数
+        :param rank: 当前进程的 ``global_rank``
+        :return: 自身
         """
         assert num_replicas<=self.num_samples, f"The number of replicas({num_replicas}) should be lesser than the " \
                                                 f"number of samples({self.num_samples})."
@@ -97,28 +97,30 @@ class UnrepeatedRandomSampler(UnrepeatedSampler):
     @property
     def num_samples(self):
         """
-        返回样本的总数
-
-        :return:
+        样本的总数
         """
         return getattr(self.dataset, 'total_len', len(self.dataset))
 
 
 class UnrepeatedSortedSampler(UnrepeatedRandomSampler):
     """
-    将 dataset 中的数据根据 length 从长到短进行迭代，并且保证在多卡场景下数据不重复。本 sampler 可能导致各个机器上的
-    batch 数量不完全一致。
+    将 ``dataset`` 中的数据根据 ``length`` 从长到短进行迭代，并且保证在多卡场景下数据不重复。
+    
+    .. note::
+    
+        本 Sampler 可能导致各个机器上的batch 数量不完全一致。
 
-    :param dataset: 实现了 __len__ 方法的数据容器。
-    :param length: 每条数据的长度。
+    :param dataset: 实现了 __len__ 方法的数据容器
+    :param length: 每条数据的长度
 
         * 为 ``List[int]`` 时
          应当与 dataset 有一样的长度，表示 dataset 中每个元素的数量；
         * 为 ``str`` 时
-         仅当传入的 ``dataset`` 是 :class:`~fastNLP.DataSet` 时，允许传入 `str` ，该 `str` 将被认为是 ``dataset`` 中的
+          仅当传入的 ``dataset`` 是 :class:`~fastNLP.DataSet` 时，允许传入 `str` ，该 `str` 将被认为是 ``dataset`` 中的
           ``field`` 。若 field 中的元素为 ``int``，则认为该值是 sample 的长度；若不为 ``int`` ，则尝试使用 ``len`` 方法
-          获取该 ``field`` 中每个元素的长度。
-    :param kwargs: fastNLP 保留使用
+          获取该 ``field`` 中每个元素的长度；
+
+    :param kwargs: fastNLP 内部使用的参数
     """
     def __init__(self, dataset, length:Union[str, List], **kwargs):
         kwargs['shuffle'] = False
@@ -146,9 +148,9 @@ class UnrepeatedSequentialSampler(UnrepeatedRandomSampler):
     按照顺序读取 dataset。
 
     :param dataset: 实现了 __len__ 方法的数据容器。
-    :param chunk_dist: 如果为 True ，当多卡时，将不间隔索取数据；为 False ，间隔取数据。例如，假设 dataset 有 10 个 sample ，使用
-        2 卡，如果为 True ，卡 0 拿 [0, 1, 2, 3, 4], 卡 1 拿 [5, 6, 7, 8, 9] ； 如果为 False ，则卡 0 拿 [0, 2, 4, 8, 8], 卡
-         1 拿 [1, 3, 5, 7, 9] 。
+    :param chunk_dist: 如果为 ``True`` ，当多卡时将不间隔索取数据；为 ``False`` 时则会间隔取数据。假设 dataset 有 10 个 sample ，使用
+        2 卡，如果为 ``True`` ，卡 **0** 拿 [0, 1, 2, 3, 4], 卡 **1** 拿 [5, 6, 7, 8, 9] ； 如果为 ``False`` ，则卡 **0** 拿 [0, 2, 4, 8, 8], 
+        卡 **1** 拿 [1, 3, 5, 7, 9] 。
     :param kwargs:
     """
     def __init__(self, dataset, chunk_dist=False, **kwargs):

fastNLP/core/utils/cache_results.py

@@ -183,9 +183,9 @@ def cal_fn_hash_code(fn: Optional[Callable] = None, fn_kwargs: Optional[dict] =
     return hasher.hexdigest()
 
 
-def cache_results(_cache_fp, _hash_param=True, _refresh=False, _verbose=1, _check_hash=True):
+def cache_results(_cache_fp: str, _hash_param: bool = True, _refresh: bool = False, _verbose: int = 1, _check_hash: bool = True):
     r"""
-    cache_results是fastNLP中用于cache数据的装饰器。通过下面的例子看一下如何使用::
+    :func:`cache_results` 是 **fastNLP** 中用于缓存数据的装饰器。通过下面的例子看一下如何使用::
 
         import time
         import numpy as np
@@ -220,19 +220,20 @@ def cache_results(_cache_fp, _hash_param=True, _refresh=False, _verbose=1, _chec
         # res = [1 8 2 5 1]
         # 2.0086121559143066
 
-    可以看到第二次运行的时候，只用了0.0001s左右，是由于第二次运行将直接从cache.pkl这个文件读取数据，而不会经过再次预处理。
-    如果在函数加上了装饰器@cache_results()，则函数会增加五个参数[_cache_fp, _hash_param, _refresh, _verbose,
-    _check_hash]。上面的例子即为使用_cache_fp的情况，这五个参数不会传入到被装饰函数中，当然被装饰函数参数名也不能包含这五个名称。
+    可以看到第二次运行的时候，只用了 0.0001s 左右，这是由于第二次运行将直接从cache.pkl这个文件读取数据，而不会经过再次预处理。
+    如果在函数加上了装饰器 ``@cache_results()``，则函数会增加五个参数 ``[_cache_fp, _hash_param, _refresh, _verbose,
+    _check_hash]``。上面的例子即为使用_cache_fp的情况，这五个参数不会传入到被装饰函数中，当然被装饰函数参数名也不能包含这五个名称。
 
-    :param str _cache_fp: 将返回结果缓存到什么位置;或从什么位置读取缓存。如果为None，cache_results没有任何效用，除非在
-        函数调用的时候传入 _cache_fp 这个参数。保存文件的名称会受到
-    :param bool _hash_param: 是否将传入给被装饰函数的 parameter 进行 str 之后的 hash 结果加入到 _cache_fp 中，这样每次函数的
+    :param _cache_fp: 将返回结果缓存到什么位置;或从什么位置读取缓存。如果为 ``None`` ，cache_results 没有任何效用，除非在
+        函数调用的时候传入 _cache_fp 这个参数。实际保存的文件名会受到 ``_hash_param`` 参数的影响，例如传入的名称是 **"caches/cache.pkl"**，
+        实际保存的文件名会是 **"caches/{hash_param_result}_cache.pkl"**。
+    :param _hash_param: 是否将传入给被装饰函数的 parameter 进行 :func:`str` 之后的 hash 结果加入到 ``_cache_fp`` 中，这样每次函数的
         parameter 改变的时候，cache 文件就自动改变了。
-    :param bool _refresh: 强制重新生成新的 cache 。
-    :param int _verbose: 是否打印cache的信息。
-    :param bool _check_hash: 如果为 True 将尝试对比修饰的函数的源码以及该函数内部调用的函数的源码的hash值。如果发现保存时的hash值
-        与当前的hash值有差异，会报warning。但该warning可能出现实质上并不影响结果的误报（例如增删空白行）；且在修改不涉及源码时，虽然
-        该修改对结果有影响，但无法做出warning。
+    :param _refresh: 强制重新生成新的 cache 。
+    :param _verbose: 是否打印 cache 的信息。
+    :param _check_hash: 如果为 ``True`` 将尝试对比修饰的函数的源码以及该函数内部调用的函数的源码的 hash 值。如果发现保存时的 hash 值
+        与当前的 hash 值有差异，会报 warning 。但该 warning 可能出现实质上并不影响结果的误报（例如增删空白行）；且在修改不涉及源码时，虽然
+        该修改对结果有影响，但无法做出 warning。
     :return:
     """
 

fastNLP/core/utils/exceptions.py

@@ -1,7 +1,7 @@
 
 class EarlyStopException(BaseException):
     r"""
-    用于EarlyStop时从Trainer训练循环中跳出。
+    用于 EarlyStop 时从 Trainer 训练循环中跳出。
 
     """
 

fastNLP/core/utils/jittor_utils.py

@@ -16,10 +16,10 @@ from fastNLP.core.dataset import Instance
 
 def is_jittor_module(model) -> bool:
     """
-    判断传入的 ``model`` 是否是 :class:`jittor.Module` 类型
+    判断传入的 ``model`` 是否是 :class:`jittor.Module` 类型。
 
-    :param model: 模型；
-    :return: 当前模型是否为 ``jittor`` 的模型；
+    :param model:
+    :return: 当前模型是否为 ``jittor`` 的模型
     """
     try:
         return isinstance(model, jt.Module)
@@ -28,10 +28,10 @@ def is_jittor_module(model) -> bool:
 
 def is_jittor_dataset(dataset) -> bool:
     """
-    判断传入的 ``dataset`` 是否是 :class:`jittor.dataset.Dataset` 类型
+    判断传入的 ``dataset`` 是否是 :class:`jittor.dataset.Dataset` 类型。
 
-    :param dataset: 数据集；
-    :return: 当前 ``dataset`` 是否为 ``jittor`` 的数据集类型；
+    :param dataset:
+    :return: 当前 ``dataset`` 是否为 ``jittor`` 的数据集类型
     """
     try:
         if isinstance(dataset, jt.dataset.Dataset):
@@ -44,7 +44,7 @@ def is_jittor_dataset(dataset) -> bool:
 
 def jittor_collate_wraps(func, auto_collator: Callable):
     """
-    对 ``jittor`` 的 ``collate_fn`` 进行 ``wrap`` 封装,。如果数据集为 ``mapping`` 类型，那么采用 ``auto_collator`` ，
+    对 ``jittor`` 的 ``collate_fn`` 进行 wrap 封装,。如果数据集为 :class:`Mapping` 类型，那么采用 ``auto_collator`` ，
     否则还是采用 ``jittor`` 的 ``collate_batch``。
 
     :param func:

fastNLP/core/utils/oneflow_utils.py

@@ -7,7 +7,7 @@ if _NEED_IMPORT_ONEFLOW:
     import oneflow
 
 __all__ = [
-    'get_oneflow_device'
+    'get_oneflow_device',
     'oneflow_move_data_to_device',
     'is_oneflow_module',
     'is_in_oneflow_dist',
@@ -32,11 +32,11 @@ def get_oneflow_device(device):
 
 def oneflow_move_data_to_device(batch: Any, device: Optional[Union[str, "oneflow.device"]] = None) -> Any:
     r"""
-    在 **oneflow** 中将数据集合 ``batch`` 传输到给定设备。任何定义方法 ``to(device)`` 的对象都将被移动并且集合中的所有其他对象将保持不变；
+    在 **oneflow** 中将数据集合 ``batch`` 传输到给定设备。
 
-    :param batch: 需要迁移的数据；
-    :param device: 数据应当迁移到的设备；当该参数的值为 ``None`` 时则不执行任何操作；
-    :return: 迁移到新设备上的数据集合；
+    :param batch: 需要迁移的数据
+    :param device: 数据应当迁移到的设备；当该参数的值为 ``None`` 时则不执行任何操作
+    :return: 迁移到新设备上的数据集合
     """
     if device is None:
         return batch
@@ -52,10 +52,10 @@ def oneflow_move_data_to_device(batch: Any, device: Optional[Union[str, "oneflow
 
 def is_oneflow_module(model) -> bool:
     """
-    判断传入的 ``model`` 是否是 :class:`oneflow.nn.Module` 类型
+    判断传入的 ``model`` 是否是 :class:`oneflow.nn.Module` 类型。
 
-    :param model: 模型；
-    :return: 当前模型是否为 ``oneflow`` 的模型；
+    :param model:
+    :return: 当前模型是否为 ``oneflow`` 的模型
     """
     try:
         return isinstance(model, oneflow.nn.Module)

fastNLP/core/utils/paddle_utils.py

@@ -42,8 +42,8 @@ def _convert_data_device(device: Union[str, int]) -> str:
 
         在分布式单进程仅支持单卡的情况下中，这个函数实际等同于直接转换为 ``gpu:0`` 返回。
 
-    :param device: 未转化的设备；
-    :return: 转化后的设备，格式为 ``gpu:x``；
+    :param device: 未转化的设备
+    :return: 转化后的设备，格式为 ``gpu:x``
     """
     try:
         user_visible_devices = os.getenv(USER_CUDA_VISIBLE_DEVICES)
@@ -65,8 +65,8 @@ def _convert_data_device(device: Union[str, int]) -> str:
 def paddle_to(data: "paddle.Tensor", device: Union[str, int, 'paddle.fluid.core_avx.Place',
                                                    'paddle.CPUPlace', 'paddle.CUDAPlace']) -> "paddle.Tensor":
     """
-    将 ``data`` 迁移到指定的 ``device`` 上。``paddle.Tensor`` 没有类似 ``torch.Tensor`` 的 ``to`` 函数，
-    该函数只是集成了 :func:`paddle.Tensor.cpu` 和 :func:`paddle.Tensor.cuda` 两个函数。
+    将 ``data`` 迁移到指定的 ``device`` 上。:class:`paddle.Tensor` 没有类似 :meth:`torch.Tensor.to` 的函数来迁移张量，
+    因此该函数只是集成了 :func:`paddle.Tensor.cpu` 和 :func:`paddle.Tensor.cuda` 两个函数。
 
     :param data: 要迁移的张量；
     :param device: 目标设备，可以是 ``str`` 或 ``int`` 及 **paddle** 自己的 :class:`paddle.fluid.core_avx.Place`、
@@ -98,7 +98,7 @@ def get_paddle_gpu_str(device: Union[str, int]) -> str:
         'gpu:1'
 
     :param device: 设备编号或设备名；
-    :return: 返回对应的 ``gpu:x`` 格式的设备名；
+    :return: 对应的 ``gpu:x`` 格式的设备名；
     """
     if isinstance(device, str):
         return device.replace("cuda", "gpu")

fastNLP/core/utils/rich_progress.py

@@ -1,5 +1,5 @@
 """
-该文件用于为 **fastNLP** 提供一个统一的 ``progress bar`` 管理，通过共用一个``Task`` 对象， :class:`~fastNLP.core.Trainer`
+该文件用于为 **fastNLP** 提供一个统一的 ``progress bar`` 管理，通过共用一个 ``Task`` 对象， :class:`~fastNLP.core.Trainer`
 中的 ``progress bar`` 和 :class:`~fastNLP.core.Evaluator` 中的 ``progress bar`` 才能不冲突
 """
 import sys

fastNLP/core/utils/utils.py

@@ -83,7 +83,7 @@ def auto_param_call(fn: Callable, *args, signature_fn: Optional[Callable] = None
         然后通过该函数签名提取参数值后，再传给 ``fn`` 进行实际的运算；
     :param mapping: 一个字典，用来更改其前面的字典的键值；
 
-    :return: 返回 ``fn`` 运行的结果；
+    :return:  ``fn`` 运行的结果；
     """
 
     if signature_fn is not None:
@@ -233,7 +233,7 @@ def check_user_specific_params(user_params: Dict, fn: Callable, fn_name=None):
         ``value`` 为每一个参数的值；
     :param fn: 将要被调用的函数；
     :param fn_name: 在打印提示信息是如何显示函数名
-    :return: 返回一个字典，其中为在之后调用函数 ``fn`` 时真正会被传进去的参数的值；
+    :return: 一个字典，其中为在之后调用函数 ``fn`` 时真正会被传进去的参数的值；
     """
     if fn_name is None:
         fn_name = fn.__name__
@@ -285,7 +285,7 @@ def match_and_substitute_params(mapping: Optional[Union[Callable, Dict]] = None,
 
     :param mapping: 用于转换的字典或者函数；当 ``mapping`` 是函数时，返回值必须为字典类型；
     :param data: 需要被转换的对象；
-    :return: 返回转换后的结果；
+    :return: 转换后的结果；
     """
     if mapping is None:
         return data

fastNLP/core/vocabulary.py

@@ -87,23 +87,18 @@ class Vocabulary(object):
         vocab["word"] # str to int
         vocab.to_word(5) # int to str
 
+    :param max_size: `Vocabulary` 的最大大小, 即能存储词的最大数量
+        若为 ``None`` , 则不限制大小。
+    :param min_freq: 能被记录下的词在文本中的最小出现频率, 应大于或等于 1。
+        若小于该频率, 词语将被视为 `unknown`. 若为 ``None`` , 所有文本中的词都被记录。
+    :param padding: padding的字符. 如果设置为 ``None`` ,
+        则vocabulary中不考虑padding, 也不计入词表大小，为 ``None`` 的情况多在为 label 建立 Vocabulary 的情况。
+    :param unknown: unknown的字符，所有未被记录的词在转为 :class:`int` 时将被视为 `unknown` 。
+        如果设置为 ``None`` ,则 vocabulary 中不考虑 `unknown`, 也不计入词表大小。
+        为 ``None`` 的情况多在为 labe l建立 Vocabulary 的情况
     """
-    
-    def __init__(self, max_size:int=None, min_freq:int=None, padding:str='<pad>', unknown:str='<unk>'):
-        r"""
-        :param max_size: `Vocabulary` 的最大大小, 即能存储词的最大数量
-            若为 ``None`` , 则不限制大小. Default: ``None``
-        :param min_freq: 能被记录下的词在文本中的最小出现频率, 应大于或等于 1.
-            若小于该频率, 词语将被视为 `unknown`. 若为 ``None`` , 所有文本中的词都被记录. Default: ``None``
-        :param padding: padding的字符. 如果设置为 ``None`` ,
-            则vocabulary中不考虑padding, 也不计入词表大小，为 ``None`` 的情况多在为label建立Vocabulary的情况.
-            Default: '<pad>'
-        :param unknown: unknown的字符，所有未被记录的词在转为 `int` 时将被视为unknown.
-            如果设置为 ``None`` ,则vocabulary中不考虑unknow, 也不计入词表大小.
-            为 ``None`` 的情况多在为label建立Vocabulary的情况.
-            Default: '<unk>'
 
-        """
+    def __init__(self, max_size:int=None, min_freq:int=None, padding:str='<pad>', unknown:str='<unk>'):
         self.max_size = max_size
         self.min_freq = min_freq
         self.word_count = Counter()
@@ -138,13 +133,16 @@ class Vocabulary(object):
         r"""
         依次增加序列中词在词典中的出现频率
 
-        :param word_lst: 列表形式的词语，如word_list=['I', 'am', 'a', 'Chinese']，列表中的每个词会计算出现频率并加入到词典中。
-        :param no_create_entry: 如果词语来自于非训练集建议设置为True。
-            如果为True，则不会有这个词语创建一个单独的entry，它将一直被指向unk的表示; 如果为False，则为这个词创建一个单独
-            的entry。如果这个word来自于dev或者test，一般设置为True，如果来自与train一般设置为False。以下两种情况: 如果新
-            加入一个word，且no_create_entry为True，但这个词之前已经在Vocabulary中且并不是no_create_entry的，则还是会为这
-            个词创建一个单独的vector; 如果no_create_entry为False，但这个词之前已经在Vocabulary中且并不是no_create_entry的，
-            则这个词将认为是需要创建单独的vector的。
+        :param word_lst: 列表形式的词语，如 word_list=['I', 'am', 'a', 'Chinese']，列表中的每个词会计算出现频率并加入到词典中。
+        :param no_create_entry: 如果词语来自于非训练集建议设置为 ``True`` 。
+            
+            * 如果为 ``True`` -- 则不会有这个词语创建一个单独的 entry ，它将一直被指向 ``<UNK>`` 的表示；
+            * 如果为 ``False`` -- 为这个词创建一个单独的 entry。如果这个词来自于验证集或训练集，一般设置为True，如果来自于训练集一
+              般设置为``False``；
+              
+            有以下两种情况: 如果新加入一个 word ，且 ``no_create_entry`` 为 ``True``，但这个词之前已经在 Vocabulary 中且并不是 
+            ``no_create_entry`` 的，则还是会为这个词创建一个单独的 vector ; 如果 ``no_create_entry`` 为 ``False`` ，但这个词之
+            前已经在 Vocabulary 中且并不是 ``no_create_entry的`` ，则这个词将认为是需要创建单独的 vector 的。
 
         """
         self._add_no_create_entry(word_lst, no_create_entry)
@@ -156,13 +154,16 @@ class Vocabulary(object):
         r"""
         增加一个新词在词典中的出现频率
 
-        :param word: 要添加进字典的新词, word为一个字符串
-        :param no_create_entry: 如果词语来自于非训练集建议设置为True。
-            如果为True，则不会有这个词语创建一个单独的entry，它将一直被指向unk的表示; 如果为False，则为这个词创建一个单独
-            的entry。如果这个word来自于dev或者test，一般设置为True，如果来自与train一般设置为False。以下两种情况: 如果新
-            加入一个word，且no_create_entry为True，但这个词之前已经在Vocabulary中且并不是no_create_entry的，则还是会为这
-            个词创建一个单独的vector; 如果no_create_entry为False，但这个词之前已经在Vocabulary中且并不是no_create_entry的，
-            则这个词将认为是需要创建单独的vector的。
+        :param word: 要添加进字典的新词， ``word`` 为一个字符串
+        :param no_create_entry: 如果词语来自于非训练集建议设置为 ``True`` 。
+            
+            * 如果为 ``True`` -- 则不会有这个词语创建一个单独的 entry ，它将一直被指向 ``<UNK>`` 的表示；
+            * 如果为 ``False`` -- 为这个词创建一个单独的 entry。如果这个词来自于验证集或训练集，一般设置为 ``True`` ，如果来自于训练集一
+              般设置为 ``False``；
+              
+            有以下两种情况: 如果新加入一个 word ，且 ``no_create_entry`` 为 ``True``，但这个词之前已经在 Vocabulary 中且并不是 
+            ``no_create_entry`` 的，则还是会为这个词创建一个单独的 vector ; 如果 ``no_create_entry`` 为 ``False`` ，但这个词之
+            前已经在 Vocabulary 中且并不是 ``no_create_entry的`` ，则这个词将认为是需要创建单独的 vector 的。
 
         """
         self._add_no_create_entry(word, no_create_entry)
@@ -173,9 +174,13 @@ class Vocabulary(object):
         r"""
         在新加入word时，检查_no_create_word的设置。
 
-        :param word: 要添加的新词或者是List类型的新词，如word='I'或者word=['I', 'am', 'a', 'Chinese']均可
-        :param no_create_entry: 如果词语来自于非训练集建议设置为True。如果为True，则不会有这个词语创建一个单独的entry，
-            它将一直被指向unk的表示; 如果为False，则为这个词创建一个单独的entry
+        :param word: 要添加的新词或者是 :class:`List`类型的新词，如 word='I' 或者 word=['I', 'am', 'a', 'Chinese'] 均可
+        :param no_create_entry: 如果词语来自于非训练集建议设置为 ``True`` 。
+            
+            * 如果为 ``True`` -- 则不会有这个词语创建一个单独的 entry ，它将一直被指向 ``<UNK>`` 的表示；
+            * 如果为 ``False`` -- 为这个词创建一个单独的 entry。如果这个词来自于验证集或训练集，一般设置为 ``True`` ，如果来自于训练集一
+              般设置为 ``False``；
+    
         :return:
 
         """
@@ -192,12 +197,16 @@ class Vocabulary(object):
         r"""
         增加一个新词在词典中的出现频率
 
-        :param word: 要添加进字典的新词, word为一个字符串
-        :param no_create_entry: 如果词语来自于非训练集建议设置为True。如果为True，则不会有这个词语创建一个单独的entry，
-            它将一直被指向unk的表示; 如果为False，则为这个词创建一个单独的entry。如果这个word来自于dev或者test，一般设置为True，
-            如果来自与train一般设置为False。以下两种情况: 如果新加入一个word，且no_create_entry为True，但这个词之前已经在Vocabulary
-            中且并不是no_create_entry的，则还是会为这词创建一个单独的vector; 如果no_create_entry为False，但这个词之前已经在Vocabulary
-            中且并不是no_create_entry的，则这个词将认为是需要创建单独的vector的。
+        :param word: 要添加进字典的新词， ``word`` 为一个字符串
+        :param no_create_entry: 如果词语来自于非训练集建议设置为 ``True`` 。
+            
+            * 如果为 ``True`` -- 则不会有这个词语创建一个单独的 entry ，它将一直被指向 ``<UNK>`` 的表示；
+            * 如果为 ``False`` -- 为这个词创建一个单独的 entry。如果这个词来自于验证集或训练集，一般设置为 ``True`` ，如果来自于训练集一
+              般设置为 ``False``；
+              
+            有以下两种情况: 如果新加入一个 word ，且 ``no_create_entry`` 为 ``True``，但这个词之前已经在 Vocabulary 中且并不是 
+            ``no_create_entry`` 的，则还是会为这个词创建一个单独的 vector ; 如果 ``no_create_entry`` 为 ``False`` ，但这个词之
+            前已经在 Vocabulary 中且并不是 ``no_create_entry的`` ，则这个词将认为是需要创建单独的 vector 的。
 
         """
         self.add(word, no_create_entry=no_create_entry)
@@ -207,12 +216,16 @@ class Vocabulary(object):
         r"""
         依次增加序列中词在词典中的出现频率
 
-        :param word_lst: 需要添加的新词的list序列，如word_lst=['I', 'am', 'a', 'Chinese']
-        :param no_create_entry: 如果词语来自于非训练集建议设置为True。如果为True，则不会有这个词语创建一个单独的entry，
-            它将一直被指向unk的表示; 如果为False，则为这个词创建一个单独的entry。如果这个word来自于dev或者test，一般设置为True，
-            如果来自与train一般设置为False。以下两种情况: 如果新加入一个word，且no_create_entry为True，但这个词之前已经在Vocabulary
-            中且并不是no_create_entry的，则还是会为这词创建一个单独的vector; 如果no_create_entry为False，但这个词之前已经在Vocabulary
-            中且并不是no_create_entry的，则这个词将认为是需要创建单独的vector的。
+        :param word_lst: 需要添加的新词的 list 序列，如 word_lst=['I', 'am', 'a', 'Chinese'] 。
+        :param no_create_entry: 如果词语来自于非训练集建议设置为 ``True`` 。
+            
+            * 如果为 ``True`` -- 则不会有这个词语创建一个单独的 entry ，它将一直被指向 ``<UNK>`` 的表示；
+            * 如果为 ``False`` -- 为这个词创建一个单独的 entry。如果这个词来自于验证集或训练集，一般设置为 ``True`` ，如果来自于训练集一
+              般设置为 ``False``；
+              
+            有以下两种情况: 如果新加入一个 word ，且 ``no_create_entry`` 为 ``True``，但这个词之前已经在 Vocabulary 中且并不是 
+            ``no_create_entry`` 的，则还是会为这个词创建一个单独的 vector ; 如果 ``no_create_entry`` 为 ``False`` ，但这个词之
+            前已经在 Vocabulary 中且并不是 ``no_create_entry的`` ，则这个词将认为是需要创建单独的 vector 的。
 
         """
         self.update(word_lst, no_create_entry=no_create_entry)
@@ -220,9 +233,8 @@ class Vocabulary(object):
     
     def build_vocab(self):
         r"""
-        根据已经出现的词和出现频率构建词典. 注意: 重复构建可能会改变词典的大小,
-        但已经记录在词典中的词, 不会改变对应的 `int`
-
+        根据已经出现的词和出现频率构建词典。注意：重复构建可能会改变词典的大小，
+        但已经记录在词典中的词，不会改变对应的 :class:`int`
         """
         if self._word2idx is None:
             self._word2idx = {}
@@ -295,16 +307,16 @@ class Vocabulary(object):
     @_check_build_vocab
     def index_dataset(self, *datasets, field_name:Union[List, str], new_field_name:Union[List, str, None]=None):
         r"""
-        将DataSet中对应field的词转为数字，Example::
+        将 ``DataSet`` 中对应 field 的词转为数字，例如::
 
             # remember to use `field_name`
             vocab.index_dataset(train_data, dev_data, test_data, field_name='words')
 
-        :param datasets: 其类型为:~fastNLP.core.Dataset或者List[~fastNLP.core.Dataset] 需要转index的一个或多个数据集
-        :param field_name: 需要转index的field, 若有多个 DataSet, 每个DataSet都必须有此 field.
-            目前支持 ``str`` , ``List[str]``
-        :param list,str new_field_name: 保存结果的field_name. 若为 ``None`` , 将覆盖原field.
-            Default: ``None``.
+        :param datasets: 其类型为 :class:`~fastNLP.core.dataset.DataSet` 或者 :class:`List` [ :class:`~fastNLP.core.dataset.DataSet` ]，
+            即需要处理的一个或多个数据集
+        :param field_name: 需要转为 index 的 field, 若有多个 DataSet, 每个 DataSet 都必须有此 field.
+            目前支持 :class:`str` , :class:`List` [ :class:`str` ]
+        :param new_field_name: 保存结果的 field_name。 若为 ``None`` , 将覆盖原 field。
         """
         
         def index_instance(field):
@@ -359,17 +371,18 @@ class Vocabulary(object):
             # remember to use `field_name`
             vocab.from_dataset(train_data1, train_data2, field_name='words', no_create_entry_dataset=[test_data1, test_data2])
 
-        :param 其类型为:~fastNLP.core.Dataset或者List[~fastNLP.core.Dataset] 需要转index的一个或多个数据集
-        :param field_name: 构建词典所使用的 field(s), 支持一个或多个field，若有多个 DataSet, 每个DataSet都必须有这些field.
+        :param datasets: 其类型为 :class:`~fastNLP.core.dataset.DataSet` 或者 List[:class:`~fastNLP.core.dataset.DataSet`]。
+        :param field_name: 构建词典所使用的 field(s), 支持一个或多个 field，若有多个 DataSet, 每个 DataSet 都必须有这些 field.
             目前支持的field结构: ``str`` , ``List[str]``
-        :param no_create_entry_dataset: 可以传入DataSet, List[DataSet]或者None(默认), 建议直接将非训练数据都传入到这个参数。该选项用在接下来的模型会使用pretrain
-            的embedding(包括glove, word2vec, elmo与bert)且会finetune的情况。如果仅使用来自于train的数据建立vocabulary，会导致test与dev
-            中的数据无法充分利用到来自于预训练embedding的信息，所以在建立词表的时候将test与dev考虑进来会使得最终的结果更好。
-            如果一个词出现在了train中，但是没在预训练模型中，embedding会为它用unk初始化，但它是单独的一个vector，如果
-            finetune embedding的话，这个词在更新之后可能会有更好的表示; 而如果这个词仅出现在了dev或test中，那么就不能为它们单独建立vector，
-            而应该让它指向unk这个vector的值。所以只位于no_create_entry_dataset中的token，将首先从预训练的词表中寻找它的表示，
-            如果找到了，就使用该表示; 如果没有找到，则认为该词的表示应该为unk的表示。
-        :return Vocabulary自身
+        :param no_create_entry_dataset: 可以传入 :class:`~fastNLP.core.dataset.DataSet`, :class:`List` [ :class:`~fastNLP.core.dataset.DataSet` ] 或者 
+            ``None`` （默认），建议直接将非训练数据都传入到这个参数。该选项用于接下来的模型会使用预训练的 embedding （包括 ``glove``, ``word2vec`` ,
+            ``elmo`` 与 ``bert`` ）且会 finetune 的情况。如果仅使用来自于训练集的数据建立词表，会导致测试集与验证集中的数据无法充分利用到来自于预训练
+            embedding 的信息，所以在建立词表的时候将测试集与验证集考虑进来会使得最终的结果更好。
+            如果一个词出现在了训练集中，但是没在预训练模型中， embedding 会为它用 ``<UNK>`` 初始化；但如果它是单独的一个 vector ，并且 finetune embedding
+            的话，这个词在更新之后可能会有更好的表示；而如果这个词仅出现在了验证集或者测试集中，那么就不能为它们单独建立 vector，而应该让它指向 ``<UNK>`` 这个
+            vector 的值。所以只位于 ``no_create_entry_dataset`` 中的 token 将首先从预训练的词表中寻找它的表示，如果找到了，就使用该表示; 如果没有找到，则认
+            为该词的表示应该为 ``<UNK>`` 的表示。
+        :return: Vocabulary 自身
 
         """
         if isinstance(field_name, str):
@@ -425,14 +438,14 @@ class Vocabulary(object):
     
     def to_index(self, w:str):
         r"""
-        将词转为数字. 若词不再词典中被记录, 将视为 unknown, 若 ``unknown=None`` , 将抛出 ``ValueError`` ::
+        将词转为数字。 若词不在词典中被记录, 将视为 `unknown`, 若 ``unknown=None`` , 将抛出 ``ValueError`` ::
 
             index = vocab.to_index('abc')
             # equals to
             index = vocab['abc']
 
         :param w: 需要输入的词语
-        :return 词语w对应的int类型的index
+        :return: 词语 ``w`` 对应的 :class:`int`类型的 index
         """
         return self.__getitem__(w)
     
@@ -440,7 +453,7 @@ class Vocabulary(object):
     @_check_build_vocab
     def unknown_idx(self):
         r"""
-        获得unknown 对应的数字.
+        获得 ``unknown`` 对应的数字.
         """
         if self.unknown is None:
             return None
@@ -450,7 +463,7 @@ class Vocabulary(object):
     @_check_build_vocab
     def padding_idx(self):
         r"""
-        获得padding 对应的数字
+        获得 ``padding`` 对应的数字
         """
         if self.padding is None:
             return None
@@ -461,16 +474,16 @@ class Vocabulary(object):
         r"""
         给定一个数字, 将其转为对应的词.
 
-        :param int idx: the index
-        :return str word: the word
+        :param idx:
+        :return: ``idx`` 对应的词
         """
         return self._idx2word[idx]
     
     def clear(self):
         r"""
-        删除Vocabulary中的词表数据。相当于重新初始化一下。
+        删除 :class:Vocabulary`` 中的词表数据。相当于重新初始化一下。
 
-        :return:
+        :return: 自身
         """
         self.word_count.clear()
         self._word2idx = None
@@ -481,7 +494,7 @@ class Vocabulary(object):
     
     def __getstate__(self):
         r"""
-        用来从pickle中加载data
+        用来从 pickle 中加载 data
 
         """
         len(self)  # make sure vocab has been built
@@ -492,7 +505,7 @@ class Vocabulary(object):
     
     def __setstate__(self, state):
         r"""
-        支持pickle的保存，保存到pickle的data state
+        支持 pickle 的保存，保存到 pickle 的 data state
 
         """
         self.__dict__.update(state)
@@ -507,11 +520,11 @@ class Vocabulary(object):
         for index in range(len(self._word2idx)):
             yield self.to_word(index), index
 
-    def save(self, filepath: [str, io.StringIO]):
+    def save(self, filepath: Union[str, io.StringIO]):
         r"""
-        :param filepath: Vocabulary的储存路径
-        :return:
+        保存当前词表。
 
+        :param filepath: 词表储存路径
         """
         if isinstance(filepath, io.IOBase):
             assert filepath.writable()
@@ -547,8 +560,8 @@ class Vocabulary(object):
         r"""
         从文件路径中加载数据
 
-        :param  filepath: Vocabulary的读取路径
-        :return: Vocabulary
+        :param filepath: 词表的读取路径
+        :return: 读取的 :class:`Vocabulary`
         """
         if isinstance(filepath, io.IOBase):
             assert filepath.writable()

fastNLP/embeddings/torch/char_embedding.py

@@ -1,6 +1,6 @@
 r"""
-该文件中主要包含的是character的Embedding，包括基于CNN与LSTM的character Embedding。与其它Embedding一样，这里的Embedding输入也是
-词的index而不需要使用词语中的char的index来获取表达。
+该文件中主要包含的是 character 的 Embedding ，包括基于 CNN 与 LSTM 的 character Embedding。与其它 Embedding 一样，这里的 Embedding 输入也是
+词的 index 而不需要使用词语中的 char 的 index 来获取表达。
 """
 
 __all__ = [
@@ -30,8 +30,8 @@ from ...core.vocabulary import Vocabulary
 
 class CNNCharEmbedding(TokenEmbedding):
     r"""
-    使用 ``CNN`` 生成 ``character embedding``。``CNN`` 的结构为, char_embed(x) -> Dropout(x) -> CNN(x) -> activation(x) -> pool -> fc -> Dropout.
-    不同的 ``kernel`` 大小的 ``fitler`` 结果是拼起来然后通过一层``fully connected layer,`` 然后输出``word``的表示。
+    使用 ``CNN`` 生成 ``character embedding``。``CNN`` 的结构为：char_embed(x) -> Dropout(x) -> CNN(x) -> activation(x) -> pool -> fc -> Dropout.
+    不同的 ``kernel`` 大小的 ``fitler`` 结果是拼起来然后通过一层 **全连接层** 然后输出 ``word`` 的表示。
 
     Example::
 
@@ -43,32 +43,33 @@ class CNNCharEmbedding(TokenEmbedding):
         >>> words = torch.LongTensor([[vocab.to_index(word) for word in "The whether is good .".split()]])
         >>> outputs = embed(words)
         >>> outputs.size()
-        # torch.Size([1, 5，50])
+        torch.Size([1, 5，50])
 
+    :param vocab: 词表
+    :param embed_size: 该 :class:`CNNCharEmbedding` 的输出维度大小。
+    :param char_emb_size: character 的 embed 的维度。character 是从 ``vocab`` 中生成的。
+    :param word_dropout: 按照一定概率随机将 word 设置为 ``unk_index`` ，这样可以使得 ``<UNK>`` 这个 token 得到足够的训练，
+        且会对网络有一定的 regularize 作用。
+    :param dropout: 以多大的概率 drop 分布式表示与 char embedding 的输出。
+    :param filter_nums: filter 的数量。长度需要和 ``kernel_sizes`` 一致。
+    :param kernel_sizes: kernel 的大小。
+    :param pool_method: character 的表示在合成一个表示时所使用的 pool 池化方法，支持 ``['avg', 'max']`` 。
+    :param activation: CNN 之后使用的激活方法，支持 ``['relu', 'sigmoid', 'tanh']`` 或者自定义函数。
+    :param min_char_freq: character 的最少出现次数。
+    :param pre_train_char_embed: 可以有两种方式调用预训练好的 :class:`CNNCharEmbedding` ：
+    
+            1. 传入 embedding 文件夹（文件夹下应该只有一个以 **.txt** 作为后缀的文件）或文件路径；
+            2. 传入 embedding 的名称，第二种情况将自动查看缓存中是否存在该模型，没有的话将自动下载；
+            3. 如果输入为 ``None`` 则使用 ``embedding_dim`` 的维度随机初始化一个 embedding；
+    :param requires_grad: 是否更新权重
+    :param include_word_start_end: 是否在每个 word 开始的 character 前和结束的 character 增加特殊标示符号
     """
     
     def __init__(self, vocab: Vocabulary, embed_size: int = 50, char_emb_size: int = 50, word_dropout: float = 0,
                  dropout: float = 0, filter_nums: List[int] = (40, 30, 20), kernel_sizes: List[int] = (5, 3, 1),
                  pool_method: str = 'max', activation='relu', min_char_freq: int = 2, pre_train_char_embed: str = None,
                  requires_grad:bool=True, include_word_start_end:bool=True):
-        r"""
-        
-        :param vocab: 词表
-        :param embed_size: 该CNNCharEmbedding的输出维度大小，默认值为50.
-        :param char_emb_size: character的embed的维度。character是从vocab中生成的。默认值为50.
-        :param word_dropout: 以多大的概率将一个词替换为unk。这样既可以训练unk也是一定的regularize。
-        :param dropout: 以多大的概率drop分布式表示与char embedding的输出。
-        :param filter_nums: filter的数量. 长度需要和kernels一致。默认值为[40, 30, 20].
-        :param kernel_sizes: kernel的大小. 默认值为[5, 3, 1].
-        :param pool_method: character的表示在合成一个表示时所使用的pool方法，支持'avg', 'max'.
-        :param activation: CNN之后使用的激活方法，支持'relu', 'sigmoid', 'tanh' 或者自定义函数.
-        :param min_char_freq: character的最少出现次数。默认值为2.
-        :param pre_train_char_embed: 可以有两种方式调用预训练好的character embedding：第一种是传入embedding文件夹
-            (文件夹下应该只有一个以.txt作为后缀的文件)或文件路径；第二种是传入embedding的名称，第二种情况将自动查看缓存中是否存在该模型，
-            没有的话将自动下载。如果输入为None则使用embedding_dim的维度随机初始化一个embedding.
-        :param requires_grad: 是否更新权重
-        :param include_word_start_end: 是否在每个word开始的character前和结束的character增加特殊标示符号；
-        """
+
         super(CNNCharEmbedding, self).__init__(vocab, word_dropout=word_dropout, dropout=dropout)
         
         for kernel in kernel_sizes:
@@ -128,10 +129,10 @@ class CNNCharEmbedding(TokenEmbedding):
 
     def forward(self, words):
         r"""
-        输入words的index后，生成对应的words的表示。
+        输入 ``words`` 的 index 后，生成对应的 ``words`` 的表示。
 
-        :param words: [batch_size, max_len]
-        :return: [batch_size, max_len, embed_size]
+        :param words: 形状为 ``[batch_size, max_len]``
+        :return: 形状为 ``[batch_size, max_len, embed_size]`` 的结果
         """
         words = self.drop_word(words)
         batch_size, max_len = words.size()
@@ -161,7 +162,7 @@ class CNNCharEmbedding(TokenEmbedding):
 
 class LSTMCharEmbedding(TokenEmbedding):
     r"""
-    使用 ``LSTM`` 的方式对 ``character`` 进行 ``encode``. embed(x) -> Dropout(x) -> LSTM(x) -> activation(x) -> pool -> Dropout
+    使用 ``LSTM`` 的方式对 ``character`` 进行 ``encode``。结构为：embed(x) -> Dropout(x) -> LSTM(x) -> activation(x) -> pool -> Dropout 。
 
     Example::
 
@@ -175,30 +176,32 @@ class LSTMCharEmbedding(TokenEmbedding):
         >>> outputs.size()
         >>> # torch.Size([1, 5，50])
 
+    :param vocab: 词表
+    :param embed_size: :class:`LSTMCharEmbedding` 的输出维度。
+    :param char_emb_size: character 的 embedding 的维度。
+    :param word_dropout: 按照一定概率随机将 word 设置为 ``unk_index`` ，这样可以使得 ``<UNK>`` 这个 token 得到足够的训练，
+        且会对网络有一定的 regularize 作用。
+    :param dropout: 以多大的概率 drop 分布式表示与 char embedding 的输出。
+    :param hidden_size: ``LSTM`` 的中间 hidden 的大小，如果为 ``bidirectional`` 的，hidden 会除二。
+    :param pool_method: character 的表示在合成一个表示时所使用的 pool 池化方法，支持 ``['avg', 'max']`` 。
+    :param activation: LSTM 之后使用的激活方法，支持 ``['relu', 'sigmoid', 'tanh']`` 或者自定义函数。
+    :param min_char_freq: character 的最少出现次数。
+    :param bidirectional: 是否使用双向的 LSTM 进行 encode。
+    :param pre_train_char_embed: 可以有两种方式调用预训练好的 :class:`LSTMCharEmbedding` ：
+    
+            1. 传入 embedding 文件夹（文件夹下应该只有一个以 **.txt** 作为后缀的文件）或文件路径；
+            2. 传入 embedding 的名称，第二种情况将自动查看缓存中是否存在该模型，
+               没有的话将自动下载；
+            3. 如果输入为 ``None`` 则使用 ``embedding_dim`` 的维度随机初始化一个 embedding；
+    :param requires_grad: 是否更新权重
+    :param include_word_start_end: 是否在每个 word 开始的 character 前和结束的 character 增加特殊标示符号
     """
     
     def __init__(self, vocab: Vocabulary, embed_size: int = 50, char_emb_size: int = 50, word_dropout: float = 0,
                  dropout: float = 0, hidden_size=50, pool_method: str = 'max', activation='relu',
                  min_char_freq: int = 2, bidirectional=True, pre_train_char_embed: str = None,
                  requires_grad:bool=True, include_word_start_end:bool=True):
-        r"""
-        
-        :param vocab: 词表
-        :param embed_size: LSTMCharEmbedding的输出维度。默认值为50.
-        :param char_emb_size: character的embedding的维度。默认值为50.
-        :param word_dropout: 以多大的概率将一个词替换为unk。这样既可以训练unk也是一定的regularize。
-        :param dropout: 以多大概率drop character embedding的输出以及最终的word的输出。
-        :param hidden_size: LSTM的中间hidden的大小，如果为bidirectional的，hidden会除二，默认为50.
-        :param pool_method: 支持'max', 'avg'。
-        :param activation: 激活函数，支持'relu', 'sigmoid', 'tanh', 或者自定义函数.
-        :param min_char_freq: character的最小出现次数。默认值为2.
-        :param bidirectional: 是否使用双向的LSTM进行encode。默认值为True。
-        :param pre_train_char_embed: 可以有两种方式调用预训练好的character embedding：第一种是传入embedding文件夹
-            (文件夹下应该只有一个以.txt作为后缀的文件)或文件路径；第二种是传入embedding的名称，第二种情况将自动查看缓存中是否存在该模型，
-            没有的话将自动下载。如果输入为None则使用embedding_dim的维度随机初始化一个embedding.
-        :param requires_grad: 是否更新权重
-        :param include_word_start_end: 是否在每个word开始的character前和结束的character增加特殊标示符号；
-        """
+
         super(LSTMCharEmbedding, self).__init__(vocab, word_dropout=word_dropout, dropout=dropout)
         
         assert hidden_size % 2 == 0, "Only even kernel is allowed."
@@ -256,10 +259,10 @@ class LSTMCharEmbedding(TokenEmbedding):
     
     def forward(self, words):
         r"""
-        输入words的index后，生成对应的words的表示。
+        输入 ``words`` 的 index 后，生成对应的 ``words`` 的表示。
 
-        :param words: [batch_size, max_len]
-        :return: [batch_size, max_len, embed_size]
+        :param words: 形状为 ``[batch_size, max_len]``
+        :return: 形状为 ``[batch_size, max_len, embed_size]`` 的结果
         """
         words = self.drop_word(words)
         batch_size, max_len = words.size()

fastNLP/embeddings/torch/embedding.py

@@ -1,5 +1,5 @@
 r"""
-该模块中的Embedding主要用于随机初始化的embedding(更推荐使用 :class:`fastNLP.embeddings.StaticEmbedding` )，或按照预训练权重初始化Embedding。
+该模块中的 :class:`Embedding` 主要用于随机初始化的 embedding （更推荐使用 :class:`fastNLP.embeddings.torch.StaticEmbedding` ），或按照预训练权重初始化 Embedding。
 
 """
 
@@ -36,19 +36,20 @@ class Embedding(Module):
         >>> init_embed = np.zeros((2000, 100))
         >>> embed = Embedding(init_embed)  # 使用numpy.ndarray的值作为初始化值初始化一个Embedding
 
+    :param init_embed: 支持传入 Embedding 的大小。支持以下类型：
+
+            1. 传入 tuple(int, int)，第一个 int 为 ``vocab_size``, 第二个 int  ``为embed_dim``；
+            2. 传入 :class:`Tensor`, :class:`Embedding`, :class:`numpy.ndarray` 等则直接使用该值初始化 Embedding；
+    
+    :param word_dropout: 按照一定概率随机将 word 设置为 ``unk_index`` ，这样可以使得 ``<UNK>`` 这个 token 得到足够的训练，
+        且会对网络有一定的 regularize 作用。设置该值时，必须同时设置 ``unk_index``。
+    :param dropout: 对 Embedding 的输出的 dropout。
+    :param unk_index: drop word 时替换为的 index。**fastNLP** 的 :class:`fastNLP.Vocabulary`` 的 ``unk_index`` 默认为 1。
     """
     
     def __init__(self, init_embed:Union[Tuple[int,int],'torch.FloatTensor','nn.Embedding',np.ndarray],
                  word_dropout:float=0, dropout:float=0.0, unk_index:int=None):
-        r"""
-        
-        :param init_embed: 支持传入Embedding的大小(传入tuple(int, int),
-            第一个int为vocab_zie, 第二个int为embed_dim); 或传入Tensor, Embedding, numpy.ndarray等则直接使用该值初始化Embedding;
-        :param word_dropout: 按照一定概率随机将word设置为unk_index，这样可以使得unk这个token得到足够的训练, 且会对网络有
-            一定的regularize的作用。设置该值时，必须同时设置unk_index
-        :param dropout: 对Embedding的输出的dropout。
-        :param unk_index: drop word时替换为的index。fastNLP的Vocabulary的unk_index默认为1。
-        """
+
         super(Embedding, self).__init__()
         
         self.embed = get_embeddings(init_embed)
@@ -69,10 +70,10 @@ class Embedding(Module):
         self.unk_index = unk_index
         self.word_dropout = word_dropout
     
-    def forward(self, words):
+    def forward(self, words: "torch.LongTensor") -> "torch.Tensor":
         r"""
-        :param torch.LongTensor words: [batch, seq_len]
-        :return: torch.Tensor : [batch, seq_len, embed_dim]
+        :param words: 形状为 ``[batch, seq_len]``
+        :return: 形状为 ``[batch, seq_len, embed_dim]`` 的张量
         """
         if self.word_dropout > 0 and self.training:
             mask = torch.ones_like(words).float() * self.word_dropout
@@ -102,7 +103,11 @@ class Embedding(Module):
     @property
     def requires_grad(self):
         r"""
-        Embedding的参数是否允许优化。True: 所有参数运行优化; False: 所有参数不允许优化; None: 部分允许优化、部分不允许
+        Embedding 的参数是否允许优化：
+        
+            - ``True`` -- 所有参数运行优化
+            - ``False`` -- 所有参数不允许优化
+            - ``None`` -- 部分允许优化、部分不允许
         :return:
         """
         if not isinstance(self.embed, TokenEmbedding):

fastNLP/embeddings/torch/stack_embedding.py

@@ -20,7 +20,7 @@ from .utils import _check_vocab_has_same_index
 
 class StackEmbedding(TokenEmbedding):
     r"""
-    支持将多个embedding集合成一个embedding。
+    支持将多个 embedding 集合成一个 embedding。
 
     Example::
 
@@ -31,16 +31,16 @@ class StackEmbedding(TokenEmbedding):
         >>> embed_2 = StaticEmbedding(vocab, model_dir_or_name='en-word2vec-300', requires_grad=True)
         >>> embed = StackEmbedding([embed_1, embed_2])
 
+    :param embeds: 一个由若干个 :class:`~fastNLP.embeddings.torch.embedding.TokenEmbedding` 组成的 :class:`list` ，要求
+        每一个 ``TokenEmbedding`` 的词表都保持一致
+    :param word_dropout: 按照一定概率随机将 word 设置为 ``unk_index`` ，这样可以使得 ``<UNK>`` 这个 token 得到足够的训练，
+        且会对网络有一定的 regularize 作用。不同 embedidng 会在相同的位置被设置为 ``<UNK>`` 。 如果这里设置了 dropout，则
+        组成的 embedding 就不要再设置 dropout 了。
+    :param dropout: 以多大的概率对 embedding 的表示进行 Dropout。0.1 即随机将 10% 的值置为 0。
     """
     
     def __init__(self, embeds: List[TokenEmbedding], word_dropout=0, dropout=0):
-        r"""
-        
-        :param embeds: 一个由若干个TokenEmbedding组成的list，要求每一个TokenEmbedding的词表都保持一致
-        :param word_dropout: 以多大的概率将一个词替换为unk。这样既可以训练unk也是一定的regularize。不同embedidng会在相同的位置
-            被设置为unknown。如果这里设置了dropout，则组成的embedding就不要再设置dropout了。
-        :param dropout: 以多大的概率对embedding的表示进行Dropout。0.1即随机将10%的值置为0。
-        """
+
         vocabs = []
         for embed in embeds:
             if hasattr(embed, 'get_word_vocab'):
@@ -59,9 +59,10 @@ class StackEmbedding(TokenEmbedding):
     
     def append(self, embed: TokenEmbedding):
         r"""
-        添加一个embedding到结尾。
+        添加一个 embedding 到结尾。
+
         :param embed:
-        :return:
+        :return: 自身
         """
         assert isinstance(embed, TokenEmbedding)
         _check_vocab_has_same_index(self.get_word_vocab(), embed.get_word_vocab())
@@ -71,8 +72,9 @@ class StackEmbedding(TokenEmbedding):
     
     def pop(self):
         r"""
-        弹出最后一个embed
-        :return:
+        弹出最后一个 embedding
+
+        :return: 被弹出的 embedding
         """
         embed = self.embeds.pop()
         self._embed_size -= embed.embed_size
@@ -81,17 +83,16 @@ class StackEmbedding(TokenEmbedding):
     @property
     def embed_size(self):
         r"""
-        该Embedding输出的vector的最后一维的维度。
-        :return:
+        该 Embedding 输出的 vector 的最后一维的维度。
         """
         return self._embed_size
     
     def forward(self, words):
         r"""
-        得到多个embedding的结果，并把结果按照顺序concat起来。
+        得到多个 embedding 的结果，并把结果按照顺序连接起来。
 
-        :param words: batch_size x max_len
-        :return: 返回的shape和当前这个stack embedding中embedding的组成有关
+        :param words: 形状为 ``[batch_size, max_len]``
+        :return: 形状和当前这个 :class:`StackEmbedding` 中 embedding 的组成有关
         """
         outputs = []
         words = self.drop_word(words)

fastNLP/embeddings/torch/static_embedding.py

@@ -10,7 +10,7 @@ import os
 from collections import defaultdict
 from copy import deepcopy
 import json
-from typing import Union
+from typing import Callable, Union
 
 import numpy as np
 
@@ -34,29 +34,27 @@ STATIC_EMBED_FILENAME = 'static.txt'
 
 class StaticEmbedding(TokenEmbedding):
     r"""
-    StaticEmbedding组件. 给定预训练embedding的名称或路径，根据vocab从embedding中抽取相应的数据(只会将出现在vocab中的词抽取出来，
-    如果没有找到，则会随机初始化一个值(但如果该word是被标记为no_create_entry的话，则不会单独创建一个值，而是会被指向unk的index))。
-    当前支持自动下载的预训练vector有:
-    
-    .. code::
-    
-        en: 实际为en-glove-840b-300d(常用)
-        en-glove-6b-50d: glove官方的50d向量
-        en-glove-6b-100d: glove官方的100d向量
-        en-glove-6b-200d: glove官方的200d向量
-        en-glove-6b-300d: glove官方的300d向量
-        en-glove-42b-300d: glove官方使用42B数据训练版本
-        en-glove-840b-300d:
-        en-glove-twitter-27b-25d:
-        en-glove-twitter-27b-50d:
-        en-glove-twitter-27b-100d:
-        en-glove-twitter-27b-200d:
-        en-word2vec-300d: word2vec官方发布的300d向量
-        en-fasttext-crawl: fasttext官方发布的300d英文预训练
-        cn-char-fastnlp-100d: fastNLP训练的100d的character embedding
-        cn-bi-fastnlp-100d: fastNLP训练的100d的bigram embedding
-        cn-tri-fastnlp-100d: fastNLP训练的100d的trigram embedding
-        cn-fasttext: fasttext官方发布的300d中文预训练embedding
+    ``StaticEmbedding`` 组件。给定预训练 embedding 的名称或路径，根据 ``vocab`` 从 embedding 中抽取相应的数据（只会将出现在 ``vocab`` 中的词抽取出来，
+    如果没有找到，则会随机初始化一个值；但如果该 word 是被标记为 ``no_create_entry`` 的话，则不会单独创建一个值，而是被指向 ``<UNK>`` 的 index）。
+    当前支持自动下载的预训练 vector 有:
+
+        - ``en`` -- 实际为 ``en-glove-840b-300d`` （常用）
+        - ``en-glove-6b-50d`` -- **glove** 官方的 50d 向量
+        - ``en-glove-6b-100d`` -- **glove** 官方的 100d 向量
+        - ``en-glove-6b-200d`` -- **glove** 官方的 200d 向量
+        - ``en-glove-6b-300d`` -- **glove** 官方的 300d 向量
+        - ``en-glove-42b-300d`` -- **glove** 官方使用 42B 数据训练版本
+        - ``en-glove-840b-300d``
+        - ``en-glove-twitter-27b-25d``
+        - ``en-glove-twitter-27b-50d``
+        - ``en-glove-twitter-27b-100d``
+        - ``en-glove-twitter-27b-200d``
+        - ``en-word2vec-300d`` -- **word2vec** 官方发布的 300d 向量
+        - ``en-fasttext-crawl`` -- **fasttext** 官方发布的 300d 英文预训练
+        - ``cn-char-fastnlp-100d`` -- **fastNLP** 训练的 100d 的 character embedding
+        - ``cn-bi-fastnlp-100d`` -- **fastNLP** 训练的 100d 的 bigram embedding
+        - ``cn-tri-fastnlp-100d`` -- **fastNLP** 训练的 100d 的 trigram embedding
+        - ``cn-fasttext`` -- **fasttext** 官方发布的 300d 中文预训练 embedding
 
     Example::
         
@@ -77,34 +75,34 @@ class StaticEmbedding(TokenEmbedding):
                      [ 0.5773,  0.7251, -0.3104,  0.0777,  0.4849],
                      [ 0.5773,  0.7251, -0.3104,  0.0777,  0.4849]]],
                    grad_fn=<EmbeddingBackward>)  # 每种word的输出是一致的。
+        
+    :param vocab: 词表。``StaticEmbedding`` 只会加载包含在词表中的词的词向量，在预训练向量中没找到的使用随机初始化
+    :param model_dir_or_name: 可以有两种方式调用预训练好的 :class:`StaticEmbedding` ：
+    
+            1. 传入 embedding 文件夹（文件夹下应该只有一个以 **.txt** 作为后缀的文件）或文件路径；
+            2. 传入 embedding 的名称，第二种情况将自动查看缓存中是否存在该模型，没有的话将自动下载;
+            3. 如果输入为 ``None`` 则使用 ``embedding_dim`` 的维度随机初始化一个 embedding；
+    :param embedding_dim: 随机初始化的 embedding 的维度，当该值为大于 0 的值时，将忽略 ``model_dir_or_name`` 。
+    :param requires_grad: 是否需要梯度。
+    :param init_method: 如何初始化没有找到的值。可以使用 :mod:`torch.nn.init` 中的各种方法，传入的方法应该接受一个 tensor，并
+        inplace 地修改其值。
+    :param lower: 是否将 ``vocab`` 中的词语小写后再和预训练的词表进行匹配。如果你的词表中包含大写的词语，或者就是需要单独
+        为大写的词语开辟一个 vector 表示，则将 ``lower`` 设置为 ``False``。
+    :param dropout: 以多大的概率对 embedding 的表示进行 Dropout。0.1 即随机将 10% 的值置为 0。
+    :param word_dropout: 按照一定概率随机将 word 设置为 ``unk_index`` ，这样可以使得 ``<UNK>`` 这个 token 得到足够的训练，
+        且会对网络有一定的 regularize 作用。
+    :param normalize: 是否对 vector 进行 ``normalize`` ，使得每个 vector 的 norm 为 1。
+    :param min_freq: Vocabulary 词频数小于这个数量的 word 将被指向 ``<UNK>``。
+    :kwargs:
+        * *only_train_min_freq* (*bool*) -- 仅对 train 中的词语使用 ``min_freq`` 筛选
+        * *only_norm_found_vector* (*bool*) -- 默认为 ``False``，是否仅对在预训练中找到的词语使用 ``normalize``
+        * *only_use_pretrain_word* (*bool*) -- 默认为 ``False``，仅使用出现在 pretrain 词表中的词，如果该词没有在预训练的词表中出现
+          则为 ``<UNK>`` 。如果 embedding 不需要更新建议设置为 ``True`` 。
 
     """
     
     def __init__(self, vocab: Vocabulary, model_dir_or_name: Union[str, None] = 'en', embedding_dim=-1, requires_grad: bool = True,
-                 init_method=None, lower=False, dropout=0, word_dropout=0, normalize=False, min_freq=1, **kwargs):
-        r"""
-        
-        :param Vocabulary vocab: 词表. StaticEmbedding只会加载包含在词表中的词的词向量，在预训练向量中没找到的使用随机初始化
-        :param model_dir_or_name: 可以有两种方式调用预训练好的static embedding：第一种是传入embedding文件夹(文件夹下应该只有一个
-            以.txt作为后缀的文件)或文件路径；第二种是传入embedding的名称，第二种情况将自动查看缓存中是否存在该模型，没有的话将自动下载。
-            如果输入为None则使用embedding_dim的维度随机初始化一个embedding。
-        :param embedding_dim: 随机初始化的embedding的维度，当该值为大于0的值时，将忽略model_dir_or_name。
-        :param requires_grad: 是否需要gradient. 默认为True
-        :param callable init_method: 如何初始化没有找到的值。可以使用torch.nn.init.*中各种方法, 传入的方法应该接受一个tensor，并
-            inplace地修改其值。
-        :param lower: 是否将vocab中的词语小写后再和预训练的词表进行匹配。如果你的词表中包含大写的词语，或者就是需要单独
-            为大写的词语开辟一个vector表示，则将lower设置为False。
-        :param dropout: 以多大的概率对embedding的表示进行Dropout。0.1即随机将10%的值置为0。
-        :param word_dropout: 以多大的概率将一个词替换为unk。这样既可以训练unk也是一定的regularize。
-        :param normalize: 是否对vector进行normalize，使得每个vector的norm为1。
-        :param min_freq: Vocabulary词频数小于这个数量的word将被指向unk。
-        :param kwargs:
-            * only_train_min_freq * (*bool*) -- 仅对 train 中的词语使用 ``min_freq`` 筛选;
-            * only_norm_found_vector * (*bool*) -- 默认为False, 是否仅对在预训练中找到的词语使用normalize;
-            * only_use_pretrain_word * (*bool*) -- 默认为False, 仅使用出现在pretrain词表中的词，如果该词没有在预训练的词表中出现
-             则为unk。如果embedding不需要更新建议设置为True。
-
-        """
+                 init_method: Callable = None, lower=False, dropout=0, word_dropout=0, normalize=False, min_freq=1, **kwargs):
         super(StaticEmbedding, self).__init__(vocab, word_dropout=word_dropout, dropout=dropout)
         if embedding_dim > 0:
             if model_dir_or_name:
@@ -327,12 +325,12 @@ class StaticEmbedding(TokenEmbedding):
 
             return vectors
     
-    def forward(self, words):
+    def forward(self, words: "torch.LongTensor") -> "torch.FloatTensor":
         r"""
-        传入words的index
+        传入 ``words`` 的 index
 
-        :param words: torch.LongTensor, [batch_size, max_len]
-        :return: torch.FloatTensor, [batch_size, max_len, embed_size]
+        :param words: 形状为 ``[batch, seq_len]``
+        :return: 形状为 ``[batch, seq_len, embed_dim]`` 的张量
         """
         if hasattr(self, 'words_to_words'):
             words = self.words_to_words[words]
@@ -341,14 +339,16 @@ class StaticEmbedding(TokenEmbedding):
         words = self.dropout(words)
         return words
 
-    def save(self, folder):
+    def save(self, folder: str):
         """
-        将embedding存储到folder下，之后可以通过使用load方法读取
+        将 embedding 存储到 ``folder`` 下，之后可以通过使用 :meth:`load` 方法读取
+
+        :param folder: 会在该 ``folder`` 下生成三个文件：
 
-        :param str folder: 会在该folder下生成三个文件, vocab.txt, static_embed_hyper.txt, static_embed_hyper.json.
-            其中vocab.txt可以用Vocabulary通过load读取; embedding.txt按照word2vec的方式存储，以空格的方式隔开元素,
-            第一行只有两个元素，剩下的行首先是word然后是各个维度的值; static_embed_hyper.json是StaticEmbedding的超参数
-        :return:
+                - ``vocab.txt``，可以通过 :meth:`fastNLP.core.Vocabulary.load` 读取；
+                - ``embedding.txt`` 按照 *word2vec* 的方式存储，以空格的方式隔开元素，第一行只有两个元素，剩下的行首先是
+                  word 然后是各个维度的值；
+                - ``static_embed_hyper.json``，:class:`StaticEmbedding` 的超参数；
         """
         os.makedirs(folder, exist_ok=True)
 
@@ -391,11 +391,11 @@ class StaticEmbedding(TokenEmbedding):
         logger.debug(f"StaticEmbedding has been saved to {folder}.")
 
     @classmethod
-    def load(cls, folder):
+    def load(cls, folder: str):
         """
 
-        :param str folder: 该folder下应该有以下三个文件vocab.txt, static_embed.txt, static_hyper.json
-        :return:
+        :param folder: 该 ``folder`` 下应该有以下三个文件 ``vocab.txt``, ``static_embed.txt``, ``static_hyper.json``
+        :return: 加载后的 embedding
         """
         for name in [VOCAB_FILENAME, STATIC_EMBED_FILENAME, STATIC_HYPER_FILENAME]:
             assert os.path.exists(os.path.join(folder, name)), f"{name} not found in {folder}."

fastNLP/embeddings/torch/utils.py

@@ -36,15 +36,17 @@ def _construct_char_vocab_from_vocab(vocab: Vocabulary, min_freq: int = 1, inclu
 
 def get_embeddings(init_embed, padding_idx=None):
     r"""
-    根据输入的init_embed返回Embedding对象。如果输入是tuple, 则随机初始化一个nn.Embedding; 如果输入是numpy.ndarray, 则按照ndarray
-    的值将nn.Embedding初始化; 如果输入是torch.Tensor, 则按该值初始化nn.Embedding; 如果输入是fastNLP中的embedding将不做处理
-    返回原对象。
-
-    :param init_embed: 可以是 tuple:(num_embedings, embedding_dim), 即embedding的大小和每个词的维度;也可以传入
-        nn.Embedding 对象, 此时就以传入的对象作为embedding; 传入np.ndarray也行，将使用传入的ndarray作为作为Embedding初始化;
-        传入torch.Tensor, 将使用传入的值作为Embedding初始化。
-    :param padding_idx: 当传入tuple时，padding_idx有效
-    :return nn.Embedding:  embeddings
+    根据输入的 ``init_embed`` 返回 ``Embedding`` 对象。
+
+    :param init_embed: 支持以下几种输入类型：
+
+            - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+            - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+            - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+            - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+
+    :param padding_idx: 当传入 :class:`tuple` 时，``padding_idx`` 有效
+    :return:
     """
     if isinstance(init_embed, tuple):
         res = nn.Embedding(
@@ -64,14 +66,14 @@ def get_embeddings(init_embed, padding_idx=None):
     return res
 
 
-def get_sinusoid_encoding_table(n_position, d_hid, padding_idx=None):
+def get_sinusoid_encoding_table(n_position: int, d_hid: int, padding_idx=None) -> "torch.FloatTensor":
     """
-    sinusoid的embedding，其中position的表示中，偶数维(0,2,4,...)是sin, 奇数(1,3,5...)是cos
+    sinusoid 的 embedding，其中 ``position`` 的表示中，偶数维 ``(0,2,4,...)`` 是 sin，奇数 ``(1,3,5...)`` 是 cos。
 
-    :param int n_position: 一共多少个position
+    :param int n_position: 一共多少个 position
     :param int d_hid: 多少维度，需要为偶数
     :param padding_idx:
-    :return: torch.FloatTensor, shape为n_position x d_hid
+    :return: 形状为 ``[n_position, d_hid]`` 的张量
     """
 
     def cal_angle(position, hid_idx):

fastNLP/envs/distributed.py

@@ -18,9 +18,9 @@ from fastNLP.envs.env import FASTNLP_GLOBAL_RANK, FASTNLP_NO_SYNC
 
 def is_cur_env_distributed() -> bool:
     """
-    单卡模式该函数一定返回 False；
-    注意进程 0 在多卡的训练模式下前后的值是不一样的，例如在开启多卡的 driver 之前，在进程 0 上的该函数返回 False；但是在开启后，在进程 0 上
-    的该函数返回的值是 True；多卡模式下除了进程 0 外的其它进程返回的值一定是 True；
+    判断当前是否处于分布式的环境下。单卡模式该函数一定返回 ``False``；
+    注意进程 0 在多卡的训练模式下前后的值是不一样的，例如在开启多卡的 driver 之前，在进程 0 上的该函数返回 ``False`` ；但是在开启后，在进程 0 上
+    的该函数返回的值是 ``True`` ；多卡模式下除了进程 0 外的其它进程返回的值一定是 ``True`` 。
     """
     return FASTNLP_GLOBAL_RANK in os.environ
 
@@ -50,8 +50,8 @@ def rank_zero_call(fn: Callable):
             return a+b
         rank_zero_call(add)(1, 2)
 
-    同时，该函数还会设置 FASTNLP_NO_SYNC 为 2，在这个环境下，所有的 fastNLP 内置的 barrier 接口，gather/broadcast 操作都没有任何
-    意义。
+    同时，该函数还会设置环境变量 ``FASTNLP_NO_SYNC`` 为 **2** ，在这个环境下，所有的 **fastNLP** 内置的 :meth:`barrier` 接口和 ``gather`` / ``broadcast``
+    操作都没有任何意义。
 
     :param fn: 需要包裹的可执行的函数。
     :return:
@@ -66,13 +66,12 @@ def rank_zero_call(fn: Callable):
 
 
 @contextmanager
-def fastnlp_no_sync_context(level=2):
+def fastnlp_no_sync_context(level: int = 2):
     """
-    用于让 fastNLP 的 barrier 以及 gather/broadcast等操作等同于只有 1 卡的多卡程序。如果为 1 表示 fastNLP 里的barrier 操作失效；
-    如果为 2 表示 barrier 与 gather/broadcast 都失效。
+    用于让 **fastNLP** 的 :meth:`barrier` 以及 ``gather`` / ``broadcast`` 等操作等同于只有 1 卡的多卡程序。如果为 1 表示 **fastNLP** 里的
+    :meth:`barrier` 操作失效；如果为 2 表示 :meth:`barrier` 与 ``gather`` / ``broadcast`` 都失效。
 
-    :param int level: 可选 [0, 1, 2]
-    :return:
+    :param level: 可选 ``[0, 1, 2]``
     """
     old_level = os.environ.get(FASTNLP_NO_SYNC, None)
     os.environ[FASTNLP_NO_SYNC] = f'{level}'
@@ -86,15 +85,13 @@ def fastnlp_no_sync_context(level=2):
 @contextmanager
 def all_rank_call_context():
     """
-    在多卡模式下，该环境内，会暂时地将 FASTNLP_GLOBAL_RANK 设置为 "0"，使得 rank_zero_call 函数失效，使得每个进程都会运行该函数。
+    在多卡模式下，该环境内，会暂时地将 ``FASTNLP_GLOBAL_RAN``K 设置为 **"0"** ，使得 :func:`rank_zero_call` 函数失效，使得每个进程都会运行该函数。
 
     使用方式::
 
         with all_rank_call_context():
             do_something  # all rank will do
 
-    :param fn:
-    :return:
     """
     old_fastnlp_global_rank = os.environ[FASTNLP_GLOBAL_RANK] if FASTNLP_GLOBAL_RANK in os.environ else None
     os.environ[FASTNLP_GLOBAL_RANK] = '0'
@@ -109,12 +106,11 @@ def all_rank_call_context():
 
 def rank_zero_rm(path: Optional[Union[str, Path]]):
     """
-    这个是因为在分布式文件系统中可能会发生错误，rank0下发删除成功后就运行走了，但实际的删除需要rank0的机器发送到远程文件系统再去执行，这个时候
-    在rank0那里，确实已经删除成功了，但是在远程文件系统那里这个操作还没完成，rank1读取的时候还是读取到存在这个文件；
-    该函数会保证所有进程都检测到 path 删除之后才退出，请保证不同进程上 path 是完全一样的，否则会陷入死锁状态。
+    仅在 rank 0 下删除文件的函数。普通的删除文件操作在分布式文件系统中可能会发生错误，rank 0 下发删除成功后就运行走了，但实际的删除需要 rank 0 的机器
+    发送到远程文件系统再去执行，这个时候在 rank 0 已经删除成功了，但是在远程文件系统那里这个操作还没完成，rank 1 读取的时候还是读取到存在这个文件；
+    该函数会保证所有进程都检测到 ``path`` 删除之后才退出，请保证不同进程上 ``path`` 是完全一样的，否则会陷入死锁状态。
 
     :param path:
-    :return:
     """
     if int(os.environ.get(FASTNLP_GLOBAL_RANK, 0)) == 0:
         if path is None:

fastNLP/envs/imports.py

@@ -26,3 +26,4 @@ _NEED_IMPORT_DEEPSPEED = _module_available("deepspeed") and 'torch' in need_impo
 _NEED_IMPORT_ONEFLOW = _module_available("oneflow") and 'oneflow' in need_import
 
 _TORCH_GREATER_EQUAL_1_8 = _NEED_IMPORT_TORCH and _compare_version("torch", operator.ge, "1.8.0")
+_TORCH_GREATER_EQUAL_1_12 = _NEED_IMPORT_TORCH and _compare_version("torch", operator.ge, "1.12.0")

fastNLP/envs/set_backend.py

@@ -10,6 +10,7 @@ from fastNLP.envs.utils import _module_available, get_gpu_count
 
 SUPPORT_BACKENDS = ['torch', 'paddle', 'jittor', 'oneflow']
 
+__all__ = []
 
 def _set_backend():
     """
@@ -152,11 +153,13 @@ def set_env(global_seed=None):
 def dump_fastnlp_backend(default:bool = False, backend=None):
     """
     将 fastNLP 的设置写入到 ~/.fastNLP/envs/ 文件夹下，
-        若 default 为 True，则保存的文件为 ~/.fastNLP/envs/default.json 。
-        如 default 为 False，则保存的文件为 ~/.fastNLP/envs/{CONDA_DEFAULT_ENV}.json ，当CONDA_DEFAULT_ENV这个环境变量不存在时
+
+        - 若 default 为 True，则保存的文件为 ~/.fastNLP/envs/default.json 。
+        - 如 default 为 False，则保存的文件为 ~/.fastNLP/envs/{CONDA_DEFAULT_ENV}.json ，当CONDA_DEFAULT_ENV这个环境变量不存在时
         ，报错。
+
     当 fastNLP 被 import 时，会默认尝试从 ~/.fastNLP/envs/{CONDA_DEFAULT_ENV}.json 读取配置文件，如果文件不存在，则尝试从
-     ~/.fastNLP/envs/default.json （如果有）读取环境变量。不过这些变量的优先级低于代码运行时的环境变量注入。
+    ~/.fastNLP/envs/default.json （如果有）读取环境变量。不过这些变量的优先级低于代码运行时的环境变量注入。
 
     会保存的环境变量为 FASTNLP_BACKEND 。
 

fastNLP/envs/set_env_on_import.py

@@ -5,6 +5,7 @@ import sys
 from .env import *
 import datetime
 
+__all__ = []
 
 def remove_local_rank_in_argv():
     """

fastNLP/io/loader/matching.py

@@ -164,7 +164,7 @@ class SNLILoader(JsonLoader):
         :param str paths: 传入一个目录, 将在该目录下寻找snli_1.0_train.jsonl, snli_1.0_dev.jsonl
             和snli_1.0_test.jsonl三个文件。
 
-        :return: 返回的 :class:`~fastNLP.io.DataBundle`
+        :return: :class:`~fastNLP.io.DataBundle`
         """
         _paths = {}
         if paths is None:

fastNLP/io/loader/summarization.py

@@ -1,66 +0,0 @@
-r"""
-.. todo::
-    doc
-"""
-
-__all__ = [
-    "ExtCNNDMLoader"
-]
-
-import os
-from typing import Union, Dict
-
-from ..data_bundle import DataBundle
-from ..utils import check_loader_paths
-from .json import JsonLoader
-
-
-class ExtCNNDMLoader(JsonLoader):
-    r"""
-    读取之后的DataSet中的field情况为
-
-    .. csv-table::
-       :header: "text", "summary", "label", "publication"
-
-       "['I got new tires from them and... ','...']", "['The new tires...','...']", "[0, 1]", "cnndm"
-       "['Don't waste your time.  We had two...','...']", "['Time is precious','...']", "[1]", "cnndm"
-       "['...']", "['...']", "[]", "cnndm"
-
-    """
-
-    def __init__(self, fields=None):
-        fields = fields or {"text": None, "summary": None, "label": None, "publication": None}
-        super(ExtCNNDMLoader, self).__init__(fields=fields)
-
-    def load(self, paths: Union[str, Dict[str, str]] = None):
-        r"""
-        从指定一个或多个路径中的文件中读取数据，返回 :class:`~fastNLP.io.DataBundle` 。
-
-        读取的field根据ExtCNNDMLoader初始化时传入的headers决定。
-
-        :param str paths: 传入一个目录, 将在该目录下寻找train.label.jsonl, dev.label.jsonl
-            test.label.jsonl三个文件(该目录还应该需要有一个名字为vocab的文件，在 :class:`~fastNLP.io.ExtCNNDMPipe`
-            当中需要用到)。
-
-        :return: 返回 :class:`~fastNLP.io.DataBundle`
-        """
-        if paths is None:
-            paths = self.download()
-        paths = check_loader_paths(paths)
-        if ('train' in paths) and ('test' not in paths):
-            paths['test'] = paths['train']
-            paths.pop('train')
-
-        datasets = {name: self._load(path) for name, path in paths.items()}
-        data_bundle = DataBundle(datasets=datasets)
-        return data_bundle
-
-    def download(self):
-        r"""
-        如果你使用了这个数据，请引用
-
-        https://arxiv.org/pdf/1506.03340.pdf
-        :return:
-        """
-        output_dir = self._get_dataset_path('ext-cnndm')
-        return output_dir

fastNLP/io/pipe/utils.py

@@ -71,7 +71,7 @@ def get_tokenizer(tokenize_method: str, lang='en'):
 
     :param str tokenize_method: 获取tokenzier方法
     :param str lang: 语言，当前仅支持en
-    :return: 返回tokenize函数
+    :return: tokenize函数
     """
     tokenizer_dict = {
         'spacy': None,

fastNLP/models/torch/biaffine_parser.py

@@ -1,5 +1,5 @@
 r"""
-Biaffine Dependency Parser 的 Pytorch 实现.
+**Biaffine Dependency Parser** 的 Pytorch 实现.
 """
 __all__ = [
     "BiaffineParser",
@@ -125,7 +125,7 @@ def _find_cycle(vertices, edges):
 
 class GraphParser(nn.Module):
     r"""
-    基于图的parser base class, 支持贪婪解码和最大生成树解码
+    基于图的 parser base class，支持 **贪婪解码** 和 **最大生成树解码**
     """
     
     def __init__(self):
@@ -134,12 +134,12 @@ class GraphParser(nn.Module):
     @staticmethod
     def greedy_decoder(arc_matrix, mask=None):
         r"""
-        贪心解码方式, 输入图, 输出贪心解码的parsing结果, 不保证合法的构成树
+        贪心解码方式，输入图，输出贪心解码的 parsing 结果，不保证合法地构成树。
 
-        :param arc_matrix: [batch, seq_len, seq_len] 输入图矩阵
-        :param mask: [batch, seq_len] 输入图的padding mask, 有内容的部分为 1, 否则为 0.
-            若为 ``None`` 时, 默认为全1向量. Default: ``None``
-        :return heads: [batch, seq_len] 每个元素在树中对应的head(parent)预测结果
+        :param arc_matrix: 输入图矩阵，形状为 ``[batch, seq_len, seq_len]``。
+        :param mask: 输入图的padding mask，形状为 ``[batch, seq_len]`` ， 有内容的部分为 **1** ， 否则为 **0** 。
+            若为 ``None`` ，则默认为全1向量。
+        :return: 每个元素在树中对应的 ``head(parent)`` 预测结果，形状为 ``[batch, seq_len]``。
         """
         _, seq_len, _ = arc_matrix.shape
         matrix = arc_matrix + torch.diag(arc_matrix.new(seq_len).fill_(-np.inf))
@@ -153,12 +153,12 @@ class GraphParser(nn.Module):
     @staticmethod
     def mst_decoder(arc_matrix, mask=None):
         r"""
-        用最大生成树算法, 计算parsing结果, 保证输出合法的树结构
+        用最大生成树算法，计算 parsing 结果，保证输出合法的树结构
 
-        :param arc_matrix: [batch, seq_len, seq_len] 输入图矩阵
-        :param mask: [batch, seq_len] 输入图的padding mask, 有内容的部分为 1, 否则为 0.
-            若为 ``None`` 时, 默认为全1向量. Default: ``None``
-        :return heads: [batch, seq_len] 每个元素在树中对应的head(parent)预测结果
+        :param arc_matrix: 输入图矩阵，形状为 ``[batch, seq_len, seq_len]``。
+        :param mask: 输入图的padding mask，形状为 ``[batch, seq_len]`` ， 有内容的部分为 **1** ， 否则为 **0** 。
+            若为 ``None`` ，则默认为全1向量。
+        :return: 每个元素在树中对应的 ``head(parent)`` 预测结果，形状为 ``[batch, seq_len]``。
         """
         batch_size, seq_len, _ = arc_matrix.shape
         matrix = arc_matrix.clone()
@@ -238,9 +238,27 @@ class LabelBilinear(nn.Module):
 
 class BiaffineParser(GraphParser):
     r"""
-    Biaffine Dependency Parser 实现.
-    论文参考 `Deep Biaffine Attention for Neural Dependency Parsing (Dozat and Manning, 2016) <https://arxiv.org/abs/1611.01734>`_ .
-
+    **Biaffine Dependency Parser** 实现。
+    论文参考 `Deep Biaffine Attention for Neural Dependency Parsing (Dozat and Manning, 2016) <https://arxiv.org/abs/1611.01734>`_ 。
+
+    :param embed: 单词词典，支持以下几种输入类型：
+
+            - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+            - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+            - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+            - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+
+    :param pos_vocab_size: part-of-speech 词典大小
+    :param pos_emb_dim: part-of-speech 向量维度
+    :param num_label: 边的类别个数
+    :param rnn_layers: rnn encoder 的层数
+    :param rnn_hidden_size: rnn encoder 的隐状态维度
+    :param arc_mlp_size: 边预测的 MLP 维度
+    :param label_mlp_size: 类别预测的 MLP 维度
+    :param dropout: dropout 概率
+    :param encoder: encoder 类别，可选 ``['lstm', 'var-lstm', 'transformer']``。
+    :param use_greedy_infer: 是否在 inference 时使用 :meth:`贪心算法 <GraphParser.greedy_decoder>` ，若为 ``False`` ，
+        将使用更加精确但相对缓慢的 :meth:`MST算法 <GraphParser.mst_decoder>` 。
     """
     
     def __init__(self,
@@ -255,23 +273,6 @@ class BiaffineParser(GraphParser):
                  dropout=0.3,
                  encoder='lstm',
                  use_greedy_infer=False):
-        r"""
-        
-        :param embed: 单词词典, 可以是 tuple, 包括(num_embedings, embedding_dim), 即
-            embedding的大小和每个词的维度. 也可以传入 nn.Embedding 对象,
-            此时就以传入的对象作为embedding
-        :param pos_vocab_size: part-of-speech 词典大小
-        :param pos_emb_dim: part-of-speech 向量维度
-        :param num_label: 边的类别个数
-        :param rnn_layers: rnn encoder的层数
-        :param rnn_hidden_size: rnn encoder 的隐状态维度
-        :param arc_mlp_size: 边预测的MLP维度
-        :param label_mlp_size: 类别预测的MLP维度
-        :param dropout: dropout概率.
-        :param encoder: encoder类别, 可选 ('lstm', 'var-lstm', 'transformer'). Default: lstm
-        :param use_greedy_infer: 是否在inference时使用贪心算法.
-            若 ``False`` , 使用更加精确但相对缓慢的MST算法. Default: ``False``
-        """
         super(BiaffineParser, self).__init__()
         rnn_out_size = 2 * rnn_hidden_size
         word_hid_dim = pos_hid_dim = rnn_hidden_size
@@ -336,20 +337,19 @@ class BiaffineParser(GraphParser):
                     nn.init.normal_(p, 0, 0.1)
     
     def forward(self, words1, words2, seq_len, target1=None):
-        r"""模型forward阶段
+        r"""
+        模型 forward 阶段
 
-        :param words1: [batch_size, seq_len] 输入word序列
-        :param words2: [batch_size, seq_len] 输入pos序列
-        :param seq_len: [batch_size, seq_len] 输入序列长度
-        :param target1: [batch_size, seq_len] 输入真实标注的heads, 仅在训练阶段有效,
-            用于训练label分类器. 若为 ``None`` , 使用预测的heads输入到label分类器
-            Default: ``None``
-        :return dict: parsing
-                结果::
+        :param words1: 输入 word 序列，形状为 ``[batch_size, seq_len]``
+        :param words2: 输入 pos 序列，形状为 ``[batch_size, seq_len]`` 
+        :param seq_len: 输入序列长度，形状为 ``[batch_size, seq_len]`` 
+        :param target1: 输入真实标注的 heads ，形状为 ``[batch_size, seq_len]`` ，仅在训练阶段有效，
+            用于训练 label 分类器. 若为 ``None`` ，则使用预测的 heads 输入到 label 分类器。
+        :return: 类型为字典的 parsing 结果，各个键的含义为：
 
-                    pred1: [batch_size, seq_len, seq_len] 边预测logits
-                    pred2: [batch_size, seq_len, num_label] label预测logits
-                    pred3: [batch_size, seq_len] heads的预测结果, 在 ``target1=None`` 时预测
+                * ``pred1`` -- **边** 预测 logits，形状为 ``[batch_size, seq_len, seq_len]``；
+                * ``pred2`` -- **label** 预测 logits，形状为 ``[batch_size, seq_len, num_label]`` ；
+                * ``pred3`` -- **heads** 的预测结果，形状为 ``[batch_size, seq_len]`` ，在 ``target1=None`` 时预测；
 
         """
         # prepare embeddings
@@ -416,6 +416,17 @@ class BiaffineParser(GraphParser):
         return res_dict
 
     def train_step(self, words1, words2, seq_len, target1, target2):
+        """
+        模型的训练接口。
+
+        :param words1: 输入 word 序列，形状为 ``[batch_size, seq_len]``
+        :param words2: 输入 pos 序列，形状为 ``[batch_size, seq_len]`` 
+        :param target1: 输入真实标注的 heads ，形状为 ``[batch_size, seq_len]`` ，仅在训练阶段有效，
+            用于训练 label 分类器. 若为 ``None`` ，则使用预测的 heads 输入到 label 分类器。
+        :param target2: 真实类别的标注，形状为 ``[batch_size, seq_len]``
+        :param seq_len: 输入序列长度，形状为 ``[batch_size, seq_len]``
+        :return: 类型为字典的结果，仅包含一个键 ``loss``，表示当次训练的 loss
+        """
         res = self(words1, words2, seq_len, target1)
         arc_pred = res['pred1']
         label_pred = res['pred2']
@@ -425,14 +436,14 @@ class BiaffineParser(GraphParser):
     @staticmethod
     def loss(pred1, pred2, target1, target2, seq_len):
         r"""
-        计算parser的loss
-
-        :param pred1: [batch_size, seq_len, seq_len] 边预测logits
-        :param pred2: [batch_size, seq_len, num_label] label预测logits
-        :param target1: [batch_size, seq_len] 真实边的标注
-        :param target2: [batch_size, seq_len] 真实类别的标注
-        :param seq_len: [batch_size, seq_len] 真实目标的长度
-        :return loss: scalar
+        计算 parser 的 loss
+
+        :param pred1: 边预测 logits，形状为 ``[batch_size, seq_len, seq_len]``
+        :param pred2: **label** 预测 logits，形状为 ``[batch_size, seq_len, num_label]``
+        :param target1: 真实边的标注，形状为 ``[batch_size, seq_len]``
+        :param target2: 真实类别的标注，形状为 ``[batch_size, seq_len]``
+        :param seq_len: 真实目标的长度，形状为 ``[batch_size, seq_len]``
+        :return: 计算出的 loss。
         """
         
         batch_size, length, _ = pred1.shape
@@ -456,14 +467,13 @@ class BiaffineParser(GraphParser):
     def evaluate_step(self, words1, words2, seq_len):
         r"""模型预测API
 
-        :param words1: [batch_size, seq_len] 输入word序列
-        :param words2: [batch_size, seq_len] 输入pos序列
-        :param seq_len: [batch_size, seq_len] 输入序列长度
-        :return dict: parsing
-                结果::
+        :param words1: 输入 word 序列，形状为 ``[batch_size, seq_len]``
+        :param words2: 输入 pos 序列，形状为 ``[batch_size, seq_len]`` 
+        :param seq_len: 输入序列长度，形状为 ``[batch_size, seq_len]``
+        :return: 字典类型的 parsing 结果，各个键的含义为：
 
-                    pred1: [batch_size, seq_len] heads的预测结果
-                    pred2: [batch_size, seq_len, num_label] label预测logits
+                * ``pred1`` -- **heads** 的预测结果，形状为 ``[batch_size, seq_len]``；
+                * ``pred2`` -- **label** 预测 logits，形状为 ``[batch_size, seq_len, num_label]`` ；
 
         """
         res = self(words1, words2, seq_len)

fastNLP/models/torch/cnn_text_classification.py

@@ -7,6 +7,7 @@ __all__ = [
     "CNNText"
 ]
 
+from typing import Union, Tuple
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -18,24 +19,27 @@ from ...modules.torch import encoder
 
 class CNNText(torch.nn.Module):
     r"""
-    使用CNN进行文本分类的模型
-    'Yoon Kim. 2014. Convolution Neural Networks for Sentence Classification.'
-    
-    """
+    使用 **CNN** 进行文本分类的模型。
+    论文参考 `Yoon Kim. 2014. Convolution Neural Networks for Sentence Classification <https://arxiv.org/abs/1408.5882>`_ 。
+
+    :param embed: 单词词典，支持以下几种输入类型：
 
+            - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+            - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+            - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+            - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+
+    :param num_classes: 一共有多少类
+    :param kernel_nums: 输出 channel 的 kernel 数目。
+        如果为 :class:`list` 或 :class:`tuple`，则需要与 ``kernel_sizes`` 的大小保持一致。
+    :param kernel_sizes: 输出 channel 的 kernel 大小。
+    :param dropout: Dropout 的大小
+    """
     def __init__(self, embed,
-                 num_classes,
-                 kernel_nums=(30, 40, 50),
-                 kernel_sizes=(1, 3, 5),
-                 dropout=0.5):
-        r"""
-        
-        :param tuple(int,int),torch.FloatTensor,nn.Embedding,numpy.ndarray embed: Embedding的大小(传入tuple(int, int),
-            第一个int为vocab_zie, 第二个int为embed_dim); 如果为Tensor, Embedding, ndarray等则直接使用该值初始化Embedding
-        :param int num_classes: 一共有多少类
-        :param int,tuple(int) kernel_sizes: 输出channel的kernel大小。
-        :param float dropout: Dropout的大小
-        """
+                 num_classes: int,
+                 kernel_nums: Union[int, Tuple[int]] = (30, 40, 50),
+                 kernel_sizes: Union[int, Tuple[int]] = (1, 3, 5),
+                 dropout: float = 0.5):
         super(CNNText, self).__init__()
 
         # no support for pre-trained embedding currently
@@ -47,14 +51,12 @@ class CNNText(torch.nn.Module):
         self.dropout = nn.Dropout(dropout)
         self.fc = nn.Linear(sum(kernel_nums), num_classes)
 
-    def forward(self, words, seq_len=None):
+    def forward(self, words: "torch.LongTensor", seq_len: "torch.LongTensor"=None):
         r"""
 
-        :param torch.LongTensor words: [batch_size, seq_len]，句子中word的index
-        :param torch.LongTensor seq_len:  [batch,] 每个句子的长度
-        :param target: 每个 sample 的目标值。
-
-        :return output:
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param seq_len:  每个句子的长度，形状为 ``[batch,]`` 
+        :return: 前向传播的结果，为仅包含一个键 ``pred`` 的字典
         """
         x = self.embed(words)  # [N,L] -> [N,L,C]
         if seq_len is not None:
@@ -70,22 +72,22 @@ class CNNText(torch.nn.Module):
     def train_step(self, words, target, seq_len=None):
         """
 
-        :param words:
-        :param target:
-        :param seq_len:
-        :return:
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param target: 每个 sample 的目标值
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]`` 
+        :return: 类型为字典的结果，仅包含一个键 ``loss``，表示当次训练的 loss
         """
         res = self(words, seq_len)
         x = res['pred']
         loss = F.cross_entropy(x, target)
         return {'loss': loss}
 
-    def evaluate_step(self, words, seq_len=None):
+    def evaluate_step(self, words: "torch.LongTensor", seq_len: "torch.LongTensor"=None):
         r"""
-        :param torch.LongTensor words: [batch_size, seq_len]，句子中word的index
-        :param torch.LongTensor seq_len:  [batch,] 每个句子的长度
 
-        :return predict: dict of torch.LongTensor, [batch_size, ]
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param seq_len:  每个句子的长度，形状为 ``[batch_size,]`` 
+        :return: 预测结果，仅包含一个键 ``pred``，值为形状为 ``[batch_size,]`` 的 :class:`torch.LongTensor`
         """
         output = self(words, seq_len)
         _, predict = output['pred'].max(dim=1)

fastNLP/models/torch/seq2seq_generator.py

@@ -1,5 +1,3 @@
-r"""undocumented"""
-
 import torch
 from torch import nn
 import torch.nn.functional as F
@@ -13,30 +11,27 @@ __all__ = ['SequenceGeneratorModel']
 
 class SequenceGeneratorModel(nn.Module):
     """
-    通过使用本模型封装seq2seq_model使得其既可以用于训练也可以用于生成。训练的时候，本模型的forward函数会被调用，生成的时候本模型的predict
-        函数会被调用。
+    通过使用本模型封装 seq2seq_model 使得其既可以用于训练也可以用于生成。训练的时候，本模型的 :meth:`forward` 函数会被调用，
+    生成的时候本模型的 :meth:`predict` 函数会被调用。
 
+    :param seq2seq_model: 序列到序列模型
+    :param bos_token_id: 句子开头的 token id
+    :param eos_token_id: 句子结束的 token id
+    :param max_length: 生成句子的最大长度, 每句话的 decode 长度为 ``max_length + max_len_a * src_len``
+    :param max_len_a: 每句话的 decode 长度为 ``max_length + max_len_a*src_len``。如果不为 0，需要保证 State 中包含 encoder_mask
+    :param num_beams: **beam search** 的大小
+    :param do_sample: 是否通过采样的方式生成
+    :param temperature: 只有在 do_sample 为 ``True`` 才有意义
+    :param top_k: 只从 ``top_k`` 中采样
+    :param top_p: 只从 ``top_p`` 的 token 中采样（ **nucleus sampling** ）
+    :param repetition_penalty: 多大程度上惩罚重复的 token
+    :param length_penalty: 对长度的惩罚，**小于 1** 鼓励长句，**大于 1** 鼓励短句
+    :param pad_token_id: 当某句话生成结束之后，之后生成的内容用 ``pad_token_id`` 补充
     """
 
-    def __init__(self, seq2seq_model: Seq2SeqModel, bos_token_id, eos_token_id=None, max_length=30, max_len_a=0.0,
-                 num_beams=1, do_sample=True, temperature=1.0, top_k=50, top_p=1.0,
-                 repetition_penalty=1, length_penalty=1.0, pad_token_id=0):
-        """
-
-        :param Seq2SeqModel seq2seq_model: 序列到序列模型
-        :param int,None bos_token_id: 句子开头的token id
-        :param int,None eos_token_id: 句子结束的token id
-        :param int max_length: 生成句子的最大长度, 每句话的decode长度为max_length + max_len_a*src_len
-        :param float max_len_a: 每句话的decode长度为max_length + max_len_a*src_len。 如果不为0，需要保证State中包含encoder_mask
-        :param int num_beams: beam search的大小
-        :param bool do_sample: 是否通过采样的方式生成
-        :param float temperature: 只有在do_sample为True才有意义
-        :param int top_k: 只从top_k中采样
-        :param float top_p: 只从top_p的token中采样，nucles sample
-        :param float repetition_penalty: 多大程度上惩罚重复的token
-        :param float length_penalty: 对长度的惩罚，小于1鼓励长句，大于1鼓励短剧
-        :param int pad_token_id: 当某句话生成结束之后，之后生成的内容用pad_token_id补充
-        """
+    def __init__(self, seq2seq_model: Seq2SeqModel, bos_token_id: int=None, eos_token_id: int=None, max_length: int=30,
+                 max_len_a: float=0.0, num_beams: int=1, do_sample: bool=True, temperature: float=1.0, top_k: int=50,
+                 top_p: float=1.0, repetition_penalty: float=1, length_penalty: float=1.0, pad_token_id: int=0):
         super().__init__()
         self.seq2seq_model = seq2seq_model
         self.generator = SequenceGenerator(seq2seq_model.decoder, max_length=max_length, max_len_a=max_len_a,
@@ -47,19 +42,28 @@ class SequenceGeneratorModel(nn.Module):
                                            repetition_penalty=repetition_penalty, length_penalty=length_penalty,
                                            pad_token_id=pad_token_id)
 
-    def forward(self, src_tokens, tgt_tokens, src_seq_len=None, tgt_seq_len=None):
+    def forward(self, src_tokens: "torch.LongTensor", tgt_tokens: "torch.LongTensor",
+                src_seq_len: "torch.LongTensor"=None, tgt_seq_len: "torch.LongTensor"=None):
         """
-        透传调用seq2seq_model的forward。
+        调用 seq2seq_model 的 :meth:`forward` 。
 
-        :param torch.LongTensor src_tokens: bsz x max_len
-        :param torch.LongTensor tgt_tokens: bsz x max_len'
-        :param torch.LongTensor src_seq_len: bsz
-        :param torch.LongTensor tgt_seq_len: bsz
-        :return:
+        :param src_tokens: source 的 token，形状为 ``[batch_size, max_len]``
+        :param tgt_tokens: target 的 token，形状为 ``[batch_size, max_len]``
+        :param src_seq_len: source的长度，形状为 ``[batch_size,]``
+        :param tgt_seq_len: target的长度，形状为 ``[batch_size,]``
+        :return: 字典 ``{'pred': torch.Tensor}``, 其中 ``pred`` 的形状为 ``[batch_size, max_len, vocab_size]``
         """
         return self.seq2seq_model(src_tokens, tgt_tokens, src_seq_len, tgt_seq_len)
 
-    def train_step(self, src_tokens, tgt_tokens, src_seq_len=None, tgt_seq_len=None):
+    def train_step(self, src_tokens: "torch.LongTensor", tgt_tokens: "torch.LongTensor",
+                    src_seq_len: "torch.LongTensor"=None, tgt_seq_len: "torch.LongTensor"=None):
+        """
+        :param src_tokens: source 的 token，形状为 ``[batch_size, max_len]``
+        :param tgt_tokens: target 的 token，形状为 ``[batch_size, max_len]``
+        :param src_seq_len: source的长度，形状为 ``[batch_size,]``
+        :param tgt_seq_len: target的长度，形状为 ``[batch_size,]``
+        :return: 字典 ``{'loss': torch.Tensor}``
+        """
         res = self(src_tokens, tgt_tokens, src_seq_len, tgt_seq_len)
         pred = res['pred']
         if tgt_seq_len is not None:
@@ -68,13 +72,13 @@ class SequenceGeneratorModel(nn.Module):
         loss = F.cross_entropy(pred[:, :-1].transpose(1, 2), tgt_tokens[:, 1:])
         return {'loss': loss}
 
-    def evaluate_step(self, src_tokens, src_seq_len=None):
+    def evaluate_step(self, src_tokens: "torch.LongTensor", src_seq_len: "torch.LongTensor"=None):
         """
-        给定source的内容，输出generate的内容。
+        给定 source 的内容，输出 generate 的内容。
 
-        :param torch.LongTensor src_tokens: bsz x max_len
-        :param torch.LongTensor src_seq_len: bsz
-        :return:
+        :param src_tokens: source 的 token，形状为 ``[batch_size, max_len]``
+        :param src_seq_len: source的长度，形状为 ``[batch_size,]``
+        :return: 字典 ``{'pred': torch.Tensor}`` ，表示生成结果
         """
         state = self.seq2seq_model.prepare_state(src_tokens, src_seq_len)
         result = self.generator.generate(state)

fastNLP/models/torch/seq2seq_model.py

@@ -18,31 +18,33 @@ __all__ = ['Seq2SeqModel', 'TransformerSeq2SeqModel', 'LSTMSeq2SeqModel']
 
 
 class Seq2SeqModel(nn.Module):
+    """
+    可以用于在 :class:`~fastNLP.core.controllers.Trainer` 中训练的 **Seq2Seq模型** 。正常情况下，继承了该函数之后，只需要
+    实现 classmethod ``build_model`` 即可。如果需要使用该模型进行生成，需要把该模型输入到 :class:`~fastNLP.models.torch.SequenceGeneratorModel`
+    中。在本模型中， :meth:`forward` 会把 encoder 后的结果传入到 decoder 中，并将 decoder 的输出 output 出来。
+
+    :param encoder: :class:`~fastNLP.modules.torch.encoder.Seq2SeqEncoder` 对象，需要实现对应的 :meth:`forward` 函数，接受两个参数，第一个为
+        ``[batch_size, max_len]`` 的 source tokens, 第二个为 ``[batch_size,]`` 的 source 的长度；需要返回两个 tensor： 
+        
+            - ``encoder_outputs`` : ``[batch_size, max_len, hidden_size]``
+            - ``encoder_mask`` :  ``[batch_size, max_len]``，为 **0** 的地方为 pad。
+        如果encoder的输出或者输入有变化，可以重载本模型的 :meth:`prepare_state` 函数或者 :meth:`forward` 函数。
+    :param decoder: :class:`~fastNLP.modules.torch.decoder.Seq2SeqEncoder` 对象，需要实现 :meth:`init_state` 函数，需要接受两个参数，分别为
+        上述的 ``encoder_outputs`` 和 ``encoder_mask``。若decoder需要更多输入，请重载当前模型的 :meth:`prepare_state` 或 :meth:`forward` 函数。
+    """
     def __init__(self, encoder: Seq2SeqEncoder, decoder: Seq2SeqDecoder):
-        """
-        可以用于在Trainer中训练的Seq2Seq模型。正常情况下，继承了该函数之后，只需要实现classmethod build_model即可。如果需要使用该模型
-            进行生成，需要把该模型输入到 :class:`~fastNLP.models.SequenceGeneratorModel` 中。在本模型中，forward()会把encoder后的
-            结果传入到decoder中，并将decoder的输出output出来。
-
-        :param encoder: Seq2SeqEncoder 对象，需要实现对应的forward()函数，接受两个参数，第一个为bsz x max_len的source tokens, 第二个为
-            bsz的source的长度；需要返回两个tensor: encoder_outputs: bsz x max_len x hidden_size, encoder_mask: bsz x max_len
-            为1的地方需要被attend。如果encoder的输出或者输入有变化，可以重载本模型的prepare_state()函数或者forward()函数
-        :param decoder: Seq2SeqDecoder 对象，需要实现init_state()函数，输出为两个参数，第一个为bsz x max_len x hidden_size是
-            encoder的输出; 第二个为bsz x max_len，为encoder输出的mask，为0的地方为pad。若decoder需要更多输入，请重载当前模型的
-            prepare_state()或forward()函数
-        """
         super().__init__()
         self.encoder = encoder
         self.decoder = decoder
 
-    def forward(self, src_tokens, tgt_tokens, src_seq_len=None, tgt_seq_len=None):
+    def forward(self, src_tokens: "torch.LongTensor", tgt_tokens: "torch.LongTensor",
+                src_seq_len: "torch.LongTensor"=None, tgt_seq_len: "torch.LongTensor"=None):
         """
-
-        :param torch.LongTensor src_tokens: source的token
-        :param torch.LongTensor tgt_tokens: target的token
-        :param torch.LongTensor src_seq_len: src的长度
-        :param torch.LongTensor tgt_seq_len: target的长度，默认用不上
-        :return: {'pred': torch.Tensor}, 其中pred的shape为bsz x max_len x vocab_size
+        :param src_tokens: source 的 token，形状为 ``[batch_size, max_len]``
+        :param tgt_tokens: target 的 token，形状为 ``[batch_size, max_len]``
+        :param src_seq_len: source的长度，形状为 ``[batch_size,]``
+        :param tgt_seq_len: target的长度，形状为 ``[batch_size,]``
+        :return: 字典 ``{'pred': torch.Tensor}``, 其中 ``pred`` 的形状为 ``[batch_size, max_len, vocab_size]``
         """
         state = self.prepare_state(src_tokens, src_seq_len)
         decoder_output = self.decoder(tgt_tokens, state)
@@ -53,7 +55,15 @@ class Seq2SeqModel(nn.Module):
         else:
             raise TypeError(f"Unsupported return type from Decoder:{type(self.decoder)}")
 
-    def train_step(self, src_tokens, tgt_tokens, src_seq_len=None, tgt_seq_len=None):
+    def train_step(self, src_tokens: "torch.LongTensor", tgt_tokens: "torch.LongTensor",
+                    src_seq_len: "torch.LongTensor"=None, tgt_seq_len: "torch.LongTensor"=None):
+        """
+        :param src_tokens: source 的 token，形状为 ``[batch_size, max_len]``
+        :param tgt_tokens: target 的 token，形状为 ``[batch_size, max_len]``
+        :param src_seq_len: source的长度，形状为 ``[batch_size,]``
+        :param tgt_seq_len: target的长度，形状为 ``[batch_size,]``
+        :return: 字典 ``{'loss': torch.Tensor}``
+        """
         res = self(src_tokens, tgt_tokens, src_seq_len, tgt_seq_len)
         pred = res['pred']
         if tgt_seq_len is not None:
@@ -62,13 +72,13 @@ class Seq2SeqModel(nn.Module):
         loss = F.cross_entropy(pred[:, :-1].transpose(1, 2), tgt_tokens[:, 1:])
         return {'loss': loss}
 
-    def prepare_state(self, src_tokens, src_seq_len=None):
+    def prepare_state(self, src_tokens: "torch.LongTensor", src_seq_len: "torch.LongTensor"=None):
         """
-        调用encoder获取state，会把encoder的encoder_output, encoder_mask直接传入到decoder.init_state中初始化一个state
+        调用 encoder 获取 state，会把 encoder 的 ``encoder_output``, ``encoder_mask`` 直接传入到 :meth:`decoder.init_state` 中初始化一个 state
 
-        :param src_tokens:
-        :param src_seq_len:
-        :return:
+        :param src_tokens: source 的 token，形状为 ``[batch_size, max_len]``
+        :param src_seq_len: source的长度，形状为 ``[batch_size,]``
+        :return: decode 初始化的 state
         """
         encoder_output, encoder_mask = self.encoder(src_tokens, src_seq_len)
         state = self.decoder.init_state(encoder_output, encoder_mask)
@@ -77,43 +87,54 @@ class Seq2SeqModel(nn.Module):
     @classmethod
     def build_model(cls, *args, **kwargs):
         """
-        需要实现本方法来进行Seq2SeqModel的初始化
+        需要实现本方法来进行 :class:`Seq2SeqModel` 的初始化
 
         :return:
         """
-        raise NotImplemented
+        raise NotImplementedError("A `Seq2SeqModel` must implement its own classmethod `build_model()`.")
 
 
 class TransformerSeq2SeqModel(Seq2SeqModel):
     """
-    Encoder为TransformerSeq2SeqEncoder, decoder为TransformerSeq2SeqDecoder，通过build_model方法初始化
-
+    Encoder 为 :class:`~fastNLP.modules.torch.encoder.TransformerSeq2SeqEncoder` ，decoder 为 
+    :class:`~fastNLP.modules.torch.decoder.TransformerSeq2SeqDecoder` 的 :class:`Seq2SeqModel` ，
+    通过 :meth:`build_model` 方法初始化。
     """
-
     def __init__(self, encoder, decoder):
         super().__init__(encoder, decoder)
 
     @classmethod
     def build_model(cls, src_embed, tgt_embed=None,
-                    pos_embed='sin', max_position=1024, num_layers=6, d_model=512, n_head=8, dim_ff=2048, dropout=0.1,
-                    bind_encoder_decoder_embed=False,
-                    bind_decoder_input_output_embed=True):
+                    pos_embed: str='sin', max_position: int=1024, num_layers: int=6, d_model: int=512,
+                    n_head: int=8, dim_ff: int=2048, dropout: float=0.1,
+                    bind_encoder_decoder_embed: bool=False,
+                    bind_decoder_input_output_embed: bool=True):
         """
-        初始化一个TransformerSeq2SeqModel
-
-        :param nn.Module, StaticEmbedding, Tuple[int, int] src_embed: source的embedding
-        :param nn.Module, StaticEmbedding, Tuple[int, int] tgt_embed: target的embedding，如果bind_encoder_decoder_embed为
-            True，则不要输入该值
-        :param str pos_embed: 支持sin, learned两种
-        :param int max_position: 最大支持长度
-        :param int num_layers: encoder和decoder的层数
-        :param int d_model: encoder和decoder输入输出的大小
-        :param int n_head: encoder和decoder的head的数量
-        :param int dim_ff: encoder和decoder中FFN中间映射的维度
-        :param float dropout: Attention和FFN dropout的大小
-        :param bool bind_encoder_decoder_embed: 是否对encoder和decoder使用相同的embedding
-        :param bool bind_decoder_input_output_embed: decoder的输出embedding是否与其输入embedding是一样的权重
-        :return: TransformerSeq2SeqModel
+        初始化一个 :class:`TransformerSeq2SeqModel` 。
+
+        :param src_embed: source 的 embedding，支持以下几种输入类型：
+
+                - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+                - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+                - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+                - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+        :param tgt_embed: target 的 embedding，如果 ``bind_encoder_decoder_embed``
+            为 ``True`` ，则不要输入该值。支持以下几种输入类型：
+
+                - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+                - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+                - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+                - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+        :param pos_embed: 支持 ``['sin', 'learned']`` 两种
+        :param max_position: 最大支持长度
+        :param num_layers: ``encoder`` 和 ``decoder`` 的层数
+        :param d_model: ``encoder`` 和 ``decoder`` 输入输出的大小
+        :param n_head: ``encoder`` 和 ``decoder`` 的 head 的数量
+        :param dim_ff: ``encoder`` 和 ``decoder`` 中 FFN 中间映射的维度
+        :param dropout: Attention 和 FFN dropout的大小
+        :param bind_encoder_decoder_embed: 是否对 ``encoder`` 和 ``decoder`` 使用相同的 embedding
+        :param bind_decoder_input_output_embed: ``decoder`` 的输出 embedding 是否与其输入 embedding 是一样的权重
+        :return: :class:`TransformerSeq2SeqModel` 模型
         """
         if bind_encoder_decoder_embed and tgt_embed is not None:
             raise RuntimeError("If you set `bind_encoder_decoder_embed=True`, please do not provide `tgt_embed`.")
@@ -152,7 +173,8 @@ class TransformerSeq2SeqModel(Seq2SeqModel):
 
 class LSTMSeq2SeqModel(Seq2SeqModel):
     """
-    使用LSTMSeq2SeqEncoder和LSTMSeq2SeqDecoder的model
+    使用 :class:`~fastNLP.modules.torch.encoder.LSTMSeq2SeqEncoder` 和 :class:`~fastNLP.modules.torch.decoder.LSTMSeq2SeqDecoder` 的
+    :class:`Seq2SeqModel`，通过 :meth:`build_model` 方法初始化。
 
     """
     def __init__(self, encoder, decoder):
@@ -160,22 +182,32 @@ class LSTMSeq2SeqModel(Seq2SeqModel):
 
     @classmethod
     def build_model(cls, src_embed, tgt_embed=None,
-                    num_layers = 3, hidden_size = 400, dropout = 0.3, bidirectional=True,
-                    attention=True, bind_encoder_decoder_embed=False,
-                    bind_decoder_input_output_embed=True):
+                    num_layers: int = 3, hidden_size: int = 400, dropout: float = 0.3, bidirectional: bool=True,
+                    attention: bool=True, bind_encoder_decoder_embed: bool=False,
+                    bind_decoder_input_output_embed: bool=True):
         """
 
-        :param nn.Module, StaticEmbedding, Tuple[int, int] src_embed: source的embedding
-        :param nn.Module, StaticEmbedding, Tuple[int, int] tgt_embed: target的embedding，如果bind_encoder_decoder_embed为
-            True，则不要输入该值
-        :param int num_layers: Encoder和Decoder的层数
-        :param int hidden_size: encoder和decoder的隐藏层大小
-        :param float dropout: 每层之间的Dropout的大小
-        :param bool bidirectional: encoder是否使用双向LSTM
-        :param bool attention: decoder是否使用attention attend encoder在所有时刻的状态
-        :param bool bind_encoder_decoder_embed: 是否对encoder和decoder使用相同的embedding
-        :param bool bind_decoder_input_output_embed: decoder的输出embedding是否与其输入embedding是一样的权重
-        :return: LSTMSeq2SeqModel
+        :param src_embed: source 的 embedding，支持以下几种输入类型：
+
+                - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+                - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+                - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+                - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+        :param tgt_embed: target 的 embedding，如果 ``bind_encoder_decoder_embed``
+            为 ``True`` ，则不要输入该值，支持以下几种输入类型：
+
+                - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+                - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+                - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+                - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+        :param num_layers: ``encoder`` 和 ``decoder`` 的层数
+        :param hidden_size: ``encoder`` 和 ``decoder`` 的隐藏层大小
+        :param dropout: 每层之间的 Dropout 的大小
+        :param bidirectional: ``encoder`` 是否使用 **双向LSTM**
+        :param attention: 是否在 ``decoder`` 中使用 attention 来添加 ``encoder`` 在所有时刻的状态
+        :param bind_encoder_decoder_embed: 是否对 ``encoder`` 和 ``decoder`` 使用相同的 embedding
+        :param bind_decoder_input_output_embed: ``decoder`` 的输出 embedding 是否与其输入 embedding 是一样的权重
+        :return: :class:`LSTMSeq2SeqModel` 模型
         """
         if bind_encoder_decoder_embed and tgt_embed is not None:
             raise RuntimeError("If you set `bind_encoder_decoder_embed=True`, please do not provide `tgt_embed`.")

fastNLP/models/torch/sequence_labeling.py

@@ -1,5 +1,5 @@
 r"""
-本模块实现了几种序列标注模型
+本模块实现了几种序列标注模型。
 """
 __all__ = [
     "SeqLabeling",
@@ -21,20 +21,24 @@ from ...modules.torch.decoder.crf import allowed_transitions
 
 class BiLSTMCRF(nn.Module):
     r"""
-    结构为embedding + BiLSTM + FC + Dropout + CRF.
+    结构为 ``Embedding`` + :class:`BiLSTM <fastNLP.modules.torch.encoder.LSTM>` + ``FC`` + ``Dropout`` +
+    :class:`CRF <fastNLP.modules.torch.decoder.ConditionalRandomField>` 。
 
+    :param embed: 支持以下几种输入类型：
+
+        - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+        - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+        - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+        - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+
+    :param num_classes: 一共多少个类
+    :param num_layers: **BiLSTM** 的层数
+    :param hidden_size: **BiLSTM** 的 ``hidden_size``，实际 hidden size 为该值的 **两倍** （前向、后向）
+    :param dropout: dropout 的概率，0 为不 dropout
+    :param target_vocab: :class:`~fastNLP.core.Vocabulary` 对象，target 与 index 的对应关系。如果传入该值，将自动避免非法的解码序列。
     """
     def __init__(self, embed, num_classes, num_layers=1, hidden_size=100, dropout=0.5,
                   target_vocab=None):
-        r"""
-        
-        :param embed: 支持(1)fastNLP的各种Embedding, (2) tuple, 指明num_embedding, dimension, 如(1000, 100)
-        :param num_classes: 一共多少个类
-        :param num_layers: BiLSTM的层数
-        :param hidden_size: BiLSTM的hidden_size，实际hidden size为该值的两倍(前向、后向)
-        :param dropout: dropout的概率，0为不dropout
-        :param target_vocab: Vocabulary对象，target与index的对应关系。如果传入该值，将自动避免非法的解码序列。
-        """
         super().__init__()
         self.embed = get_embeddings(embed)
 
@@ -55,7 +59,13 @@ class BiLSTMCRF(nn.Module):
 
         self.crf = ConditionalRandomField(num_classes, include_start_end_trans=True, allowed_transitions=trans)
 
-    def forward(self, words, seq_len=None, target=None):
+    def forward(self, words: "torch.LongTensor", target: "torch.LongTensor"=None, seq_len: "torch.LongTensor"=None):
+        """
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param target: 每个 sample 的目标值
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]``
+        :return: 如果 ``target`` 为 ``None``，则返回预测结果 ``{'pred': torch.Tensor}``，否则返回 loss ``{'loss': torch.Tensor}``
+        """
         words = self.embed(words)
         feats, _ = self.lstm(words, seq_len=seq_len)
         feats = self.fc(feats)
@@ -69,28 +79,40 @@ class BiLSTMCRF(nn.Module):
             loss = self.crf(logits, target, mask).mean()
             return {'loss':loss}
 
-    def train_step(self, words, seq_len, target):
+    def train_step(self, words: "torch.LongTensor", target: "torch.LongTensor", seq_len: "torch.LongTensor"):
+        """
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param target: 每个 sample 的目标值
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]``
+        :return: 如果 ``target`` 为 ``None``，则返回预测结果 ``{'pred': torch.Tensor}``，否则返回 loss ``{'loss': torch.Tensor}``
+        """
         return self(words, seq_len, target)
 
-    def evaluate_step(self, words, seq_len):
+    def evaluate_step(self, words: "torch.LongTensor", seq_len: "torch.LongTensor"):
+        """
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]``
+        :return: 预测结果 ``{'pred': torch.Tensor}``
+        """
         return self(words, seq_len)
 
 
 class SeqLabeling(nn.Module):
     r"""
-    一个基础的Sequence labeling的模型。
-    用于做sequence labeling的基础类。结构包含一层Embedding，一层LSTM(单向，一层)，一层FC，以及一层CRF。
-    
+    一个基础的 Sequence labeling 的模型。
+    用于做 sequence labeling 的基础类。结构包含一层 ``Embedding`` ，一层 :class:`~fastNLP.modules.torch.encoder.LSTM` (单向，一层)，
+    一层全连接层，以及一层 :class:`CRF <fastNLP.modules.torch.decoder.ConditionalRandomField>` 。
+
+    :param embed: 支持以下几种输入类型：
+
+            - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+            - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+            - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+            - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+    :param hidden_size: :class:`fastNLP.modules.torch.encoder.LSTM` 隐藏层的大小
+    :param num_classes: 一共有多少类
     """
-    
-    def __init__(self, embed, hidden_size, num_classes):
-        r"""
-        
-        :param tuple(int,int),torch.FloatTensor,nn.Embedding,numpy.ndarray embed: Embedding的大小(传入tuple(int, int),
-            第一个int为vocab_zie, 第二个int为embed_dim); 如果为Tensor, embedding, ndarray等则直接使用该值初始化Embedding
-        :param int hidden_size: LSTM隐藏层的大小
-        :param int num_classes: 一共有多少类
-        """
+    def __init__(self, embed, hidden_size: int, num_classes: int):
         super(SeqLabeling, self).__init__()
         
         self.embedding = get_embeddings(embed)
@@ -98,11 +120,11 @@ class SeqLabeling(nn.Module):
         self.fc = nn.Linear(hidden_size, num_classes)
         self.crf = decoder.ConditionalRandomField(num_classes)
 
-    def forward(self, words, seq_len):
+    def forward(self, words: "torch.LongTensor", seq_len: "torch.LongTensor"):
         r"""
-        :param torch.LongTensor words: [batch_size, max_len]，序列的index
-        :param torch.LongTensor seq_len: [batch_size,], 这个序列的长度
-        :return
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]``
+        :return: 预测结果 ``{'pred': torch.Tensor}``
         """
         x = self.embedding(words)
         # [batch_size, max_len, word_emb_dim]
@@ -112,19 +134,23 @@ class SeqLabeling(nn.Module):
         return {'pred': x}
         # [batch_size, max_len, num_classes]
 
-    def train_step(self, words, seq_len, target):
+    def train_step(self, words, target, seq_len):
+        """
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param target: 每个 sample 的目标值
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]``
+        :return: 如果 ``target`` 为 ``None``，则返回预测结果 ``{'pred': torch.Tensor}``，否则返回 loss ``{'loss': torch.Tensor}``
+        """
         res = self(words, seq_len)
         pred = res['pred']
         mask = seq_len_to_mask(seq_len, max_len=target.size(1))
         return {'loss': self._internal_loss(pred, target, mask)}
 
     def evaluate_step(self, words, seq_len):
-        r"""
-        用于在预测时使用
-
-        :param torch.LongTensor words: [batch_size, max_len]
-        :param torch.LongTensor seq_len: [batch_size,]
-        :return: {'pred': xx}, [batch_size, max_len]
+        """
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]``
+        :return: 预测结果 ``{'pred': torch.Tensor}``
         """
         mask = seq_len_to_mask(seq_len, max_len=words.size(1))
 
@@ -158,22 +184,25 @@ class SeqLabeling(nn.Module):
 
 class AdvSeqLabel(nn.Module):
     r"""
-    更复杂的Sequence Labelling模型。结构为Embedding, LayerNorm, 双向LSTM(两层)，FC，LayerNorm，DropOut，FC，CRF。
+    更复杂的 Sequence Labelling 模型。结构为 ``Embedding``, ``LayerNorm``, :class:`BiLSTM <fastNLP.modules.torch.encoder.LSTM>` （两层），
+    ``FC``，``LayerNorm``，``Dropout``，``FC``，:class:`CRF <fastNLP.modules.torch.decoder.ConditionalRandomField>`。
+
+    :param embed: 支持以下几种输入类型：
+
+        - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+        - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+        - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+        - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+    :param hidden_size: :class:`~fastNLP.modules.torch.LSTM` 的隐藏层大小
+    :param num_classes: 有多少个类
+    :param dropout: :class:`~fastNLP.modules.torch.LSTM` 中以及 DropOut 层的 drop 概率
+    :param id2words: tag id 转为其 tag word 的表。用于在 CRF 解码时防止解出非法的顺序，比如 **'BMES'** 这个标签规范中，``'S'``
+        不能出现在 ``'B'`` 之后。这里也支持类似于 ``'B-NN'``，即 ``'-'`` 前为标签类型的指示，后面为具体的 tag 的情况。这里不但会保证
+        ``'B-NN'`` 后面不为 ``'S-NN'`` 还会保证 ``'B-NN'`` 后面不会出现 ``'M-xx'`` （任何非 ``'M-NN'`` 和 ``'E-NN'`` 的情况）。
+    :param encoding_type: 支持 ``["BIO", "BMES", "BEMSO"]`` ，只有在 ``id2words`` 不为 ``None`` 的情况有用。
     """
     
-    def __init__(self, embed, hidden_size, num_classes, dropout=0.3, id2words=None, encoding_type='bmes'):
-        r"""
-        
-        :param tuple(int,int),torch.FloatTensor,nn.Embedding,numpy.ndarray embed: Embedding的大小(传入tuple(int, int),
-            第一个int为vocab_zie, 第二个int为embed_dim); 如果为Tensor, Embedding, ndarray等则直接使用该值初始化Embedding
-        :param int hidden_size: LSTM的隐层大小
-        :param int num_classes: 有多少个类
-        :param float dropout: LSTM中以及DropOut层的drop概率
-        :param dict id2words: tag id转为其tag word的表。用于在CRF解码时防止解出非法的顺序，比如'BMES'这个标签规范中，'S'
-            不能出现在'B'之后。这里也支持类似与'B-NN'，即'-'前为标签类型的指示，后面为具体的tag的情况。这里不但会保证
-            'B-NN'后面不为'S-NN'还会保证'B-NN'后面不会出现'M-xx'(任何非'M-NN'和'E-NN'的情况。)
-        :param str encoding_type: 支持"BIO", "BMES", "BEMSO", 只有在id2words不为None的情况有用。
-        """
+    def __init__(self, embed, hidden_size: int, num_classes: int, dropout: float=0.3, id2words: dict=None, encoding_type: str='bmes'):
         super().__init__()
         
         self.Embedding = get_embeddings(embed)
@@ -217,13 +246,12 @@ class AdvSeqLabel(nn.Module):
         total_loss = self.Crf(x, y, mask)
         return torch.mean(total_loss)
     
-    def forward(self, words, seq_len, target=None):
-        r"""
-        :param torch.LongTensor words: [batch_size, mex_len]
-        :param torch.LongTensor seq_len:[batch_size, ]
-        :param torch.LongTensor target: [batch_size, max_len]
-        :return y: If truth is None, return list of [decode path(list)]. Used in testing and predicting.
-                   If truth is not None, return loss, a scalar. Used in training.
+    def forward(self, words: "torch.LongTensor", target: "torch.LongTensor"=None, seq_len: "torch.LongTensor"=None):
+        """
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param target: 每个 sample 的目标值
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]``
+        :return: 如果 ``target`` 为 ``None``，则返回预测结果 ``{'pred': torch.Tensor}``，否则返回 loss ``{'loss': torch.Tensor}``
         """
         
         words = words.long()
@@ -251,21 +279,19 @@ class AdvSeqLabel(nn.Module):
         else:
             return {"pred": self._decode(x, mask)}
     
-    def train_step(self, words, seq_len, target):
-        r"""
-        
-        :param torch.LongTensor words: [batch_size, mex_len]
-        :param torch.LongTensor seq_len: [batch_size, ]
-        :param torch.LongTensor target: [batch_size, max_len], 目标
-        :return torch.Tensor: a scalar loss
+    def train_step(self, words: "torch.LongTensor", target: "torch.LongTensor", seq_len: "torch.LongTensor"):
+        """
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param target: 每个 sample 的目标值
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]``
+        :return: 如果 ``target`` 为 ``None``，则返回预测结果 ``{'pred': torch.Tensor}``，否则返回 loss ``{'loss': torch.Tensor}``
         """
         return self(words, seq_len, target)
     
-    def evaluate_step(self, words, seq_len):
-        r"""
-        
-        :param torch.LongTensor words: [batch_size, mex_len]
-        :param torch.LongTensor seq_len: [batch_size, ]
-        :return torch.LongTensor: [batch_size, max_len]
+    def evaluate_step(self, words: "torch.LongTensor", seq_len: "torch.LongTensor"):
+        """
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]``
+        :return: 预测结果 ``{'pred': torch.Tensor}``
         """
         return self(words, seq_len)

fastNLP/modules/torch/attention.py

@@ -1,5 +1,3 @@
-r"""undocumented"""
-
 __all__ = [
     "MultiHeadAttention",
     "BiAttention",
@@ -17,7 +15,11 @@ from .decoder.seq2seq_state import TransformerState
 
 class DotAttention(nn.Module):
     r"""
-    Transformer当中的DotAttention
+    **Transformer** 当中的 **DotAttention**
+
+    :param key_size:
+    :param value_size:
+    :param dropout:
     """
 
     def __init__(self, key_size, value_size, dropout=0.0):
@@ -31,10 +33,10 @@ class DotAttention(nn.Module):
     def forward(self, Q, K, V, mask_out=None):
         r"""
 
-        :param Q: [..., seq_len_q, key_size]
-        :param K: [..., seq_len_k, key_size]
-        :param V: [..., seq_len_k, value_size]
-        :param mask_out: [..., 1, seq_len] or [..., seq_len_q, seq_len_k]
+        :param Q: ``[..., seq_len_q, key_size]``
+        :param K: ``[..., seq_len_k, key_size]``
+        :param V: ``[..., seq_len_k, value_size]``
+        :param mask_out: ``[..., 1, seq_len]`` or ``[..., seq_len_q, seq_len_k]``
         """
         output = torch.matmul(Q, K.transpose(-1, -2)) / self.scale
         if mask_out is not None:
@@ -46,8 +48,12 @@ class DotAttention(nn.Module):
 
 class MultiHeadAttention(nn.Module):
     """
-    Attention is all you need中提到的多头注意力
+    `Attention is all you need <https://arxiv.org/abs/1706.03762>`_ 中提到的多头注意力
 
+    :param d_model:
+    :param n_head:
+    :param dropout:
+    :param layer_idx:
     """
     def __init__(self, d_model: int = 512, n_head: int = 8, dropout: float = 0.0, layer_idx: int = None):
         super(MultiHeadAttention, self).__init__()
@@ -69,12 +75,13 @@ class MultiHeadAttention(nn.Module):
     def forward(self, query, key, value, key_mask=None, attn_mask=None, state=None):
         """
 
-        :param query: batch x seq x dim
-        :param key: batch x seq x dim
-        :param value: batch x seq x dim
-        :param key_mask: batch x seq 用于指示哪些key不要attend到；注意到mask为1的地方是要attend到的
-        :param attn_mask: seq x seq, 用于mask掉attention map。 主要是用在训练时decoder端的self attention，下三角为1
-        :param state: 过去的信息，在inference的时候会用到，比如encoder output、decoder的prev kv。这样可以减少计算。
+        :param query: ``[batch, seq, dim]``
+        :param key: ``[batch, seq, dim]``
+        :param value: ``[batch, seq, dim]``
+        :param key_mask: ``[batch, seq]`` 用于指示哪些 ``key`` 不要 attend 到；注意到 mask 为 **1** 的地方是要attend到的
+        :param attn_mask: ``[seq, seq]``, 用于 mask 掉 attention map。 主要是用在训练时 decoder 端的 :class:`SelfAttention` ，
+            下三角为 1。
+        :param state: 过去的信息，在 inference 的时候会用到，比如 encoder output、decoder 的 prev kv。这样可以减少计算。
         :return:
         """
         assert key.size() == value.size()
@@ -149,15 +156,15 @@ class MultiHeadAttention(nn.Module):
 
 
 class AttentionLayer(nn.Module):
-    def __init__(selfu, input_size, key_dim, value_dim, bias=False):
-        """
-        可用于LSTM2LSTM的序列到序列模型的decode过程中，该attention是在decode过程中根据上一个step的hidden计算对encoder结果的attention
+    """
+    可用于 LSTM2LSTM 的序列到序列模型的 decode 过程中，该 attention 是在 decode 过程中根据上一个 step 的 hidden 计算对 encoder 结果的 attention
 
-        :param int input_size: 输入的大小
-        :param int key_dim: 一般就是encoder_output输出的维度
-        :param int value_dim: 输出的大小维度, 一般就是decoder hidden的大小
-        :param bias:
-        """
+    :param int input_size: 输入的大小
+    :param int key_dim: 一般就是 encoder_output 输出的维度
+    :param int value_dim: 输出的大小维度, 一般就是 decoder hidden 的大小
+    :param bias:
+    """
+    def __init__(selfu, input_size, key_dim, value_dim, bias=False):
         super().__init__()
 
         selfu.input_proj = nn.Linear(input_size, key_dim, bias=bias)
@@ -166,10 +173,10 @@ class AttentionLayer(nn.Module):
     def forward(self, input, encode_outputs, encode_mask):
         """
 
-        :param input: batch_size x input_size
-        :param encode_outputs: batch_size x max_len x key_dim
-        :param encode_mask: batch_size x max_len, 为0的地方为padding
-        :return: hidden: batch_size x value_dim, scores: batch_size x max_len, normalized过的
+        :param input: ``[batch_size, input_size]``
+        :param encode_outputs: ``[batch_size, max_len, key_dim]``
+        :param encode_mask: ``[batch_size, max_len]``, 为0的地方为padding
+        :return: hidden: ``[batch_size, value_dim]``, scores: ``[batch_size, max_len]``, normalized 过的
         """
 
         # x: bsz x encode_hidden_size
@@ -221,9 +228,9 @@ def _weighted_sum(tensor, weights, mask):
 
 class BiAttention(nn.Module):
     r"""
-    Bi Attention module
+    **Bi Attention module**
 
-    对于给定的两个向量序列 :math:`a_i` 和 :math:`b_j` , BiAttention模块将通过以下的公式来计算attention结果
+    对于给定的两个向量序列 :math:`a_i` 和 :math:`b_j` , :class:`BiAttention` 模块将通过以下的公式来计算 attention 结果
 
     .. math::
 
@@ -237,11 +244,14 @@ class BiAttention(nn.Module):
 
     def forward(self, premise_batch, premise_mask, hypothesis_batch, hypothesis_mask):
         r"""
-        :param torch.Tensor premise_batch: [batch_size, a_seq_len, hidden_size]
-        :param torch.Tensor premise_mask: [batch_size, a_seq_len]
-        :param torch.Tensor hypothesis_batch: [batch_size, b_seq_len, hidden_size]
-        :param torch.Tensor hypothesis_mask: [batch_size, b_seq_len]
-        :return: torch.Tensor attended_premises: [batch_size, a_seq_len, hidden_size] torch.Tensor attended_hypotheses: [batch_size, b_seq_len, hidden_size]
+        :param premise_batch: ``[batch_size, a_seq_len, hidden_size]``
+        :param premise_mask: ``[batch_size, a_seq_len]``
+        :param hypothesis_batch: ``[batch_size, b_seq_len, hidden_size]``
+        :param hypothesis_mask: ``[batch_size, b_seq_len]``
+        :return: 一个包含两个张量的元组，分别为：
+
+                - ``attended_premises`` : ``[batch_size, a_seq_len, hidden_size]``
+                - ``attended_hypotheses`` : ``[batch_size, b_seq_len, hidden_size]``
         """
         similarity_matrix = premise_batch.bmm(hypothesis_batch.transpose(2, 1)
                                               .contiguous())
@@ -264,17 +274,15 @@ class BiAttention(nn.Module):
 class SelfAttention(nn.Module):
     r"""
     这是一个基于论文 `A structured self-attentive sentence embedding <https://arxiv.org/pdf/1703.03130.pdf>`_
-    的Self Attention Module.
+    的 **Self Attention Module** 。
+
+    :param  input_size: 输入 tensor 的 hidden 维度
+    :param attention_unit: 输出 tensor 的 hidden 维度
+    :param attention_hops:
+    :param drop: dropout 概率
     """
 
     def __init__(self, input_size, attention_unit=300, attention_hops=10, drop=0.5):
-        r"""
-        
-        :param int input_size: 输入tensor的hidden维度
-        :param int attention_unit: 输出tensor的hidden维度
-        :param int attention_hops:
-        :param float drop: dropout概率，默认值为0.5
-        """
         super(SelfAttention, self).__init__()
 
         self.attention_hops = attention_hops
@@ -301,10 +309,12 @@ class SelfAttention(nn.Module):
 
     def forward(self, input, input_origin):
         r"""
-        :param torch.Tensor input: [batch_size, seq_len, hidden_size] 要做attention的矩阵
-        :param torch.Tensor input_origin: [batch_size, seq_len] 原始token的index组成的矩阵，含有pad部分内容
-        :return torch.Tensor output1: [batch_size, multi-head, hidden_size] 经过attention操作后输入矩阵的结果
-        :return torch.Tensor output2: [1] attention惩罚项，是一个标量
+        :param input: 要做 **attention** 的矩阵，形状为 ``[batch_size, seq_len, hidden_size]``
+        :param input_origin: 原始 token 的 index 组成的矩阵，含有 pad 部分内容，形状为 ``[batch_size, seq_len]`` 
+        :return: 一个元组，分别是：
+
+                - 经过 **attention** 操作后输入矩阵的结果，形状为 ``[batch_size, multi-head, hidden_size]``
+                - **attention** 惩罚项，是一个标量
         """
         input = input.contiguous()
         size = input.size()  # [bsz, len, nhid]

fastNLP/modules/torch/decoder/crf.py

@@ -1,11 +1,9 @@
-r"""undocumented"""
-
 __all__ = [
     "ConditionalRandomField",
     "allowed_transitions"
 ]
 
-from typing import Union, List
+from typing import Union, List, Tuple
 
 import torch
 from torch import nn
@@ -14,17 +12,19 @@ from ....core.metrics.span_f1_pre_rec_metric import _get_encoding_type_from_tag_
 from ....core.vocabulary import Vocabulary
 
 
-def allowed_transitions(tag_vocab:Union[Vocabulary, dict], encoding_type:str=None, include_start_end:bool=False):
+def allowed_transitions(tag_vocab:Union[Vocabulary, dict], encoding_type:str=None, include_start_end:bool=False) -> List[Tuple[int, int]]:
     r"""
-    给定一个id到label的映射表，返回所有可以跳转的(from_tag_id, to_tag_id)列表。
-
-    :param tag_vocab: 支持类型为tag或tag-label。只有tag的,比如"B", "M"; 也可以是"B-NN", "M-NN",
-        tag和label之间一定要用"-"隔开。如果传入dict，格式需要形如{0:"O", 1:"B-tag1"}，即index在前，tag在后。
-    :param encoding_type: 支持``["bio", "bmes", "bmeso", "bioes"]``。默认为None，通过vocab自动推断
-    :param include_start_end: 是否包含开始与结尾的转换。比如在bio中，b/o可以在开头，但是i不能在开头；
-        为True，返回的结果中会包含(start_idx, b_idx), (start_idx, o_idx), 但是不包含(start_idx, i_idx);
-        start_idx=len(id2label), end_idx=len(id2label)+1。为False, 返回的结果中不含与开始结尾相关的内容
-    :return: List[Tuple(int, int)]], 内部的Tuple是可以进行跳转的(from_tag_id, to_tag_id)。
+    给定一个 ``id`` 到 ``label`` 的映射表，返回所有可以跳转的 ``(from_tag_id, to_tag_id)`` 列表。
+
+    :param tag_vocab: 支持类型为 tag 或 tag-label 。只有 tag 的,比如 ``"B"`` 、 ``"M"``，也可以是 ``"B-NN"`` 、 ``"M-NN"``，
+        tag 和 label之间一定要用 ``"-"`` 隔开。如果传入 :class:`dict` ，格式需要形如 ``{0:"O", 1:"B-tag1"}`` ，即 **index 在前，tag 在后** 。
+    :param encoding_type: 支持 ``["bio", "bmes", "bmeso", "bioes", None]``。默认为 ``None``，通过 ``tag_vocab`` 自动推断
+    :param include_start_end: 是否包含开始与结尾的转换。比如在 ``bio`` 中， ``b/o`` 可以在开头，但是 ``i`` 不能在开头；
+        
+            - 为 ``True`` -- 返回的结果中会包含 ``(start_idx, b_idx), (start_idx, o_idx)`` ，但是不包含 ``(start_idx, i_idx)`` ；
+              其中 ``start_idx=len(id2label)``，``end_idx=len(id2label)+1`` ；
+            - 为 ``False`` , 返回的结果中不含与开始结尾相关的内容；
+    :return: 一系列元组构成的列表，内部的 :class:`Tuple` 是可以进行跳转的 ``(from_tag_id, to_tag_id)``。
     """
     if encoding_type is None:
         encoding_type = _get_encoding_type_from_tag_vocab(tag_vocab)
@@ -167,19 +167,15 @@ def _is_transition_allowed(encoding_type, from_tag, from_label, to_tag, to_label
 
 class ConditionalRandomField(nn.Module):
     r"""
-    条件随机场。提供 forward() 以及 viterbi_decode() 两个方法，分别用于训练与inference。
+    条件随机场。提供 :meth:`forward` 以及 :meth:`viterbi_decode` 两个方法，分别用于 **训练** 与 **inference** 。
 
+    :param num_tags: 标签的数量
+    :param include_start_end_trans: 是否考虑各个 tag 作为开始以及结尾的分数。
+    :param allowed_transitions: 内部的 ``Tuple[from_tag_id(int), to_tag_id(int)]`` 视为允许发生的跃迁，其他没
+        有包含的跃迁认为是禁止跃迁，可以通过 :func:`allowed_transitions` 函数得到；如果为 ``None`` ，则所有跃迁均为合法。
     """
 
     def __init__(self, num_tags:int, include_start_end_trans:bool=False, allowed_transitions:List=None):
-        r"""
-        
-        :param num_tags: 标签的数量
-        :param include_start_end_trans: 是否考虑各个tag作为开始以及结尾的分数。
-        :param allowed_transitions: 内部的Tuple[from_tag_id(int),
-                                   to_tag_id(int)]视为允许发生的跃迁，其他没有包含的跃迁认为是禁止跃迁，可以通过
-                                   allowed_transitions()函数得到；如果为None，则所有跃迁均为合法
-        """
         super(ConditionalRandomField, self).__init__()
 
         self.include_start_end_trans = include_start_end_trans
@@ -213,9 +209,9 @@ class ConditionalRandomField(nn.Module):
         r"""Computes the (batch_size,) denominator term for the log-likelihood, which is the
         sum of the likelihoods across all possible state sequences.
 
-        :param logits:FloatTensor, max_len x batch_size x num_tags
-        :param mask:ByteTensor, max_len x batch_size
-        :return:FloatTensor, batch_size
+        :param logits:FloatTensor, ``[max_len, batch_size, num_tags]``
+        :param mask:ByteTensor, ``[max_len, batch_size]``
+        :return:FloatTensor, ``[batch_size,]``
         """
         seq_len, batch_size, n_tags = logits.size()
         alpha = logits[0]
@@ -239,10 +235,10 @@ class ConditionalRandomField(nn.Module):
     def _gold_score(self, logits, tags, mask):
         r"""
         Compute the score for the gold path.
-        :param logits: FloatTensor, max_len x batch_size x num_tags
-        :param tags: LongTensor, max_len x batch_size
-        :param mask: ByteTensor, max_len x batch_size
-        :return:FloatTensor, batch_size
+        :param logits: FloatTensor, ``[max_len, batch_size, num_tags]``
+        :param tags: LongTensor, ``[max_len, batch_size]``
+        :param mask: ByteTensor, ``[max_len, batch_size]``
+        :return:FloatTensor, ``[batch_size.]``
         """
         seq_len, batch_size, _ = logits.size()
         batch_idx = torch.arange(batch_size, dtype=torch.long, device=logits.device)
@@ -265,14 +261,14 @@ class ConditionalRandomField(nn.Module):
         # return [B,]
         return score
 
-    def forward(self, feats, tags, mask):
+    def forward(self, feats: "torch.FloatTensor", tags: "torch.LongTensor", mask: "torch.ByteTensor") -> "torch.FloatTensor":
         r"""
-        用于计算CRF的前向loss，返回值为一个batch_size的FloatTensor，可能需要mean()求得loss。
+        用于计算 ``CRF`` 的前向 loss，返回值为一个形状为 ``[batch_size,]`` 的 :class:`torch.FloatTensor` ，可能需要 :func:`mean` 求得 loss 。
 
-        :param torch.FloatTensor feats: batch_size x max_len x num_tags，特征矩阵。
-        :param torch.LongTensor tags: batch_size x max_len，标签矩阵。
-        :param torch.ByteTensor mask: batch_size x max_len，为0的位置认为是padding。
-        :return: torch.FloatTensor, (batch_size,)
+        :param feats: 特征矩阵，形状为 ``[batch_size, max_len, num_tags]``
+        :param tags: 标签矩阵，形状为 ``[batch_size, max_len]``
+        :param mask: 形状为 ``[batch_size, max_len]`` ，为 **0** 的位置认为是 padding。
+        :return: ``[batch_size,]``
         """
         feats = feats.transpose(0, 1)
         tags = tags.transpose(0, 1).long()
@@ -282,17 +278,20 @@ class ConditionalRandomField(nn.Module):
 
         return all_path_score - gold_path_score
 
-    def viterbi_decode(self, logits, mask, unpad=False):
-        r"""给定一个特征矩阵以及转移分数矩阵，计算出最佳的路径以及对应的分数
-
-        :param torch.FloatTensor logits: batch_size x max_len x num_tags，特征矩阵。
-        :param torch.ByteTensor mask: batch_size x max_len, 为0的位置认为是pad；如果为None，则认为没有padding。
-        :param bool unpad: 是否将结果删去padding。False, 返回的是batch_size x max_len的tensor; True，返回的是
-            List[List[int]], 内部的List[int]为每个sequence的label，已经除去pad部分，即每个List[int]的长度是这
-            个sample的有效长度。
-        :return: 返回 (paths, scores)。
-                    paths: 是解码后的路径, 其值参照unpad参数.
-                    scores: torch.FloatTensor, size为(batch_size,), 对应每个最优路径的分数。
+    def viterbi_decode(self, logits: "torch.FloatTensor", mask: "torch.ByteTensor", unpad=False):
+        r"""给定一个 **特征矩阵** 以及 **转移分数矩阵** ，计算出最佳的路径以及对应的分数
+
+        :param logits: 特征矩阵，形状为 ``[batch_size, max_len, num_tags]``
+        :param mask: 标签矩阵，形状为 ``[batch_size, max_len]`` ，为 **0** 的位置认为是 padding。如果为 ``None`` ，则认为没有 padding。
+        :param unpad: 是否将结果删去 padding：
+
+                - 为 ``False`` 时，返回的是 ``[batch_size, max_len]`` 的张量
+                - 为 ``True`` 时，返回的是 :class:`List` [:class:`List` [ :class:`int` ]], 内部的 :class:`List` [:class:`int` ] 为每个 
+                  sequence 的 label ，已经除去 pad 部分，即每个 :class:`List` [ :class:`int` ] 的长度是这个 sample 的有效长度。
+        :return:  (paths, scores)。
+
+                - ``paths`` -- 解码后的路径, 其值参照 ``unpad`` 参数.
+                - ``scores`` -- :class:`torch.FloatTensor` ，形状为 ``[batch_size,]`` ，对应每个最优路径的分数。
 
         """
         batch_size, max_len, n_tags = logits.size()

fastNLP/modules/torch/decoder/mlp.py

@@ -1,23 +1,15 @@
-r"""undocumented"""
-
 __all__ = [
     "MLP"
 ]
 
+from typing import List, Callable, Union
 import torch
 import torch.nn as nn
 
 
 class MLP(nn.Module):
     r"""
-    多层感知器
-
-    
-    .. note::
-        隐藏层的激活函数通过activation定义。一个str/function或者一个str/function的list可以被传入activation。
-        如果只传入了一个str/function，那么所有隐藏层的激活函数都由这个str/function定义；
-        如果传入了一个str/function的list，那么每一个隐藏层的激活函数由这个list中对应的元素定义，其中list的长度为隐藏层数。
-        输出层的激活函数由output_activation定义，默认值为None，此时输出层没有激活函数。
+    多层感知器。
         
     Examples::
 
@@ -31,18 +23,20 @@ class MLP(nn.Module):
         >>>     y = net(x)
         >>>     print(x)
         >>>     print(y)
+
+    :param size_layer: 一个 int 的列表，用来定义 :class:`MLP` 的层数，列表中的数字为每一层是 hidden 数目。 :class:`MLP` 的层数为 ``len(size_layer) - 1``
+    :param activation: 隐藏层的激活函数，可以支持多种类型：
+    
+            - 一个 :class:`str` 或函数 -- 所有隐藏层的激活函数都为 ``activation`` 代表的函数；
+            - :class:`str` 或函数的列表 -- 每一个隐藏层的激活函数都为列表中对应的函数，其中列表长度为隐藏层数；
+
+        对于字符串类型的输入，支持 ``['relu', 'tanh', 'sigmoid']`` 三种激活函数。
+    :param output_activation: 输出层的激活函数。默认值为 ``None``，表示输出层没有激活函数
+    :param dropout: dropout 概率
     """
 
-    def __init__(self, size_layer, activation='relu', output_activation=None, initial_method=None, dropout=0.0):
-        r"""
-        
-        :param List[int] size_layer: 一个int的列表，用来定义MLP的层数，列表中的数字为每一层是hidden数目。MLP的层数为 len(size_layer) - 1
-        :param Union[str,func,List[str]] activation: 一个字符串或者函数的列表，用来定义每一个隐层的激活函数，字符串包括relu，tanh和
-            sigmoid，默认值为relu
-        :param Union[str,func] output_activation:  字符串或者函数，用来定义输出层的激活函数，默认值为None，表示输出层没有激活函数
-        :param str initial_method: 参数初始化方式
-        :param float dropout: dropout概率，默认值为0
-        """
+    def __init__(self, size_layer: List[int], activation: Union[str, Callable, List[str]]='relu',
+                output_activation: Union[str, Callable]=None, dropout: float=0.0):
         super(MLP, self).__init__()
         self.hiddens = nn.ModuleList()
         self.output = None
@@ -85,8 +79,8 @@ class MLP(nn.Module):
 
     def forward(self, x):
         r"""
-        :param torch.Tensor x: MLP接受的输入
-        :return: torch.Tensor : MLP的输出结果
+        :param x:
+        :return:
         """
         for layer, func in zip(self.hiddens, self.hidden_active):
             x = self.dropout(func(layer(x)))

fastNLP/modules/torch/decoder/seq2seq_decoder.py

@@ -1,4 +1,3 @@
-r"""undocumented"""
 from typing import Union, Tuple
 import math
 
@@ -16,54 +15,52 @@ __all__ = ['Seq2SeqDecoder', 'TransformerSeq2SeqDecoder', 'LSTMSeq2SeqDecoder']
 
 class Seq2SeqDecoder(nn.Module):
     """
-    Sequence-to-Sequence Decoder的基类。一定需要实现forward、decode函数，剩下的函数根据需要实现。每个Seq2SeqDecoder都应该有相应的State对象
-        用来承载该Decoder所需要的Encoder输出、Decoder需要记录的历史信息(例如LSTM的hidden信息)。
-
+    **Sequence-to-Sequence Decoder** 的基类。一定需要实现 :meth:`forward` 和 :meth:`decode` 函数，剩下的函数根据需要实现。每个 ``Seq2SeqDecoder`` 都应该有相应的
+    :class:`~fastNLP.modules.torch.decoder.State` 对象用来承载该 ``Decoder`` 所需要的 ``Encoder`` 输出、``Decoder`` 需要记录的历史信（例如 :class:`~fastNLP.modules.torch.encoder.LSTM`
+    的 hidden 信息）。
     """
     def __init__(self):
         super().__init__()
 
-    def forward(self, tokens, state, **kwargs):
+    def forward(self, tokens: "torch.LongTensor", state: State, **kwargs):
         """
 
-        :param torch.LongTensor tokens: bsz x max_len
-        :param State state: state包含了encoder的输出以及decode之前的内容
-        :return: 返回值可以为bsz x max_len x vocab_size的Tensor，也可以是一个list，但是第一个元素必须是词的预测分布
+        :param tokens: ``[batch_size, max_len]``
+        :param state: ``state`` 包含了 ``encoder`` 的输出以及 ``decode`` 之前的内容
+        :return: 返回值可以为 ``[batch_size, max_len, vocab_size]`` 的张量，也可以是一个 :class:`list`，但是第一个元素必须是词的预测分布
         """
         raise NotImplemented
 
-    def reorder_states(self, indices, states):
+    def reorder_states(self, indices: torch.LongTensor, states):
         """
-        根据indices重新排列states中的状态，在beam search进行生成时，会用到该函数。
+        根据 ``indices`` 重新排列 ``states`` 中的状态，在 ``beam search`` 进行生成时，会用到该函数。
 
-        :param torch.LongTensor indices:
-        :param State states:
-        :return:
+        :param indices:
+        :param states:
         """
         assert isinstance(states, State), f"`states` should be of type State instead of {type(states)}"
         states.reorder_state(indices)
 
-    def init_state(self, encoder_output, encoder_mask):
+    def init_state(self, encoder_output: Union[torch.Tensor, list, tuple], encoder_mask: Union[torch.Tensor, list, tuple]):
         """
-        初始化一个state对象，用来记录了encoder的输出以及decode已经完成的部分。
+        初始化一个 :class:`~fastNLP.modules.torch.decoder.State` 对象，用来记录 ``encoder`` 的输出以及 ``decode`` 已经完成的部分。
 
-        :param Union[torch.Tensor, list, tuple] encoder_output: 如果不为None，内部元素需要为torch.Tensor, 默认其中第一维是batch
+        :param encoder_output: 如果不为 ``None`` ，内部元素需要为 :class:`torch.Tensor`，默认其中第一维是 batch_size
             维度
-        :param Union[torch.Tensor, list, tuple] encoder_mask: 如果部位None，内部元素需要torch.Tensor, 默认其中第一维是batch
+        :param encoder_mask: 如果不为 ``None``，内部元素需要为 :class:`torch.Tensor`，默认其中第一维是 batch_size
             维度
-        :param kwargs:
-        :return: State, 返回一个State对象，记录了encoder的输出
+        :return: 一个 :class:`~fastNLP.modules.torch.decoder.State` 对象，记录了 ``encoder`` 的输出
         """
         state = State(encoder_output, encoder_mask)
         return state
 
-    def decode(self, tokens, state):
+    def decode(self, tokens: torch.LongTensor, state) -> torch.FloatTensor:
         """
-        根据states中的内容，以及tokens中的内容进行之后的生成。
+        根据 ``states`` 中的内容，以及 ``tokens`` 中的内容进行之后的生成。
 
-        :param torch.LongTensor tokens: bsz x max_len, 截止到上一个时刻所有的token输出。
-        :param State state: 记录了encoder输出与decoder过去状态
-        :return: torch.FloatTensor: bsz x vocab_size, 输出的是下一个时刻的分布
+        :param tokens: ``[batch_size, max_len]``，截止到上一个时刻所有的 token 输出。
+        :param state: 记录了 ``encoder`` 输出与 ``decoder`` 过去状态
+        :return: `下一个时刻的分布，形状为 ``[batch_size, vocab_size]``
         """
         outputs = self(state=state, tokens=tokens)
         if isinstance(outputs, torch.Tensor):
@@ -84,8 +81,8 @@ class TiedEmbedding(nn.Module):
     def forward(self, x):
         """
 
-        :param torch.FloatTensor x: bsz x * x embed_size
-        :return: torch.FloatTensor bsz x * x vocab_size
+        :param torch.FloatTensor x: batch_size x * x embed_size
+        :return: torch.FloatTensor batch_size x * x vocab_size
         """
         return torch.matmul(x, self.weight.t())
 
@@ -110,18 +107,24 @@ def get_bind_decoder_output_embed(embed):
 
 class LSTMSeq2SeqDecoder(Seq2SeqDecoder):
     """
-    LSTM的Decoder
-
-    :param nn.Module,tuple embed: decoder输入的embedding.
-    :param int num_layers: 多少层LSTM
-    :param int hidden_size: 隐藏层大小, 该值也被认为是encoder的输出维度大小
-    :param dropout: Dropout的大小
-    :param bool bind_decoder_input_output_embed: 是否将输出层和输入层的词向量绑定在一起（即为同一个），若embed为StaticEmbedding，
-        则StaticEmbedding的vocab不能包含no_create_entry的token，同时StaticEmbedding初始化时lower为False, min_freq=1.
-    :param bool attention: 是否使用attention
+    **LSTM** 的 Decoder
+
+    :param embed: ``decoder`` 输入的 embedding，支持以下几种输入类型：
+
+            - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+            - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+            - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+            - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+    :param num_layers: LSTM 的层数
+    :param hidden_size: 隐藏层大小, 该值也被认为是 ``encoder`` 的输出维度大小
+    :param dropout: Dropout 的大小
+    :param bind_decoder_input_output_embed: ``decoder`` 的输出 embedding 是否与其输入 embedding 是一样的权重（即为同一个），若 ``embed`` 为 
+        :class:`~fastNLP.embeddings.StaticEmbedding`，则 ``StaticEmbedding`` 的 ``vocab`` 不能包含 ``no_create_entry`` 的 token ，同时
+        ``StaticEmbedding`` 初始化时 ``lower`` 为 ``False``，``min_freq=1``。
+    :param attention: 是否使用attention
     """
-    def __init__(self, embed: Union[nn.Module, Tuple[int, int]], num_layers = 3, hidden_size = 300,
-                 dropout = 0.3, bind_decoder_input_output_embed = True, attention=True):
+    def __init__(self, embed: Union[nn.Module, Tuple[int, int]], num_layers: int = 3, hidden_size: int = 300,
+                 dropout: float = 0.3, bind_decoder_input_output_embed: bool = True, attention: bool = True):
         super().__init__()
         self.embed = get_embeddings(init_embed=embed)
         self.embed_dim = embed.embedding_dim
@@ -141,13 +144,14 @@ class LSTMSeq2SeqDecoder(Seq2SeqDecoder):
         self.output_proj = nn.Linear(hidden_size, self.embed_dim)
         self.dropout_layer = nn.Dropout(dropout)
 
-    def forward(self, tokens, state, return_attention=False):
+    def forward(self, tokens: torch.LongTensor, state: LSTMState, return_attention: bool=False):
         """
 
-        :param torch.LongTensor tokens: batch x max_len
-        :param LSTMState state: 保存encoder输出和decode状态的State对象
-        :param bool return_attention: 是否返回attention的的score
-        :return: bsz x max_len x vocab_size; 如果return_attention=True, 还会返回bsz x max_len x encode_length
+        :param tokens: ``[batch_size, max_len]``
+        :param state: 保存 ``encoder`` 输出和 ``decode`` 状态的 :class:`~fastNLP.modules.torch.decoder.LSTMState` 对象
+        :param return_attention: 是否返回 attention 的 score
+        :return: 形状为 ``[batch_size, max_len, vocab_size]`` 的结果。如果 ``return_attention=True`` 则返回一个元组，一个元素为结果，第二个结果为
+            注意力权重，形状为 ``[batch_size, max_len, encode_length]``
         """
         src_output = state.encoder_output
         encoder_mask = state.encoder_mask
@@ -196,14 +200,18 @@ class LSTMSeq2SeqDecoder(Seq2SeqDecoder):
             return feats, attn_weights
         return feats
 
-    def init_state(self, encoder_output, encoder_mask) -> LSTMState:
+    def init_state(self, encoder_output, encoder_mask: torch.ByteTensor) -> LSTMState:
         """
 
-        :param encoder_output: 输入可以有两种情况(1) 输入为一个tuple，包含三个内容(encoder_output, (hidden, cell))，其中encoder_output:
-            bsz x max_len x hidden_size, hidden: bsz x hidden_size, cell:bsz x hidden_size,一般使用LSTMEncoder的最后一层的
-            hidden state和cell state来赋值这两个值
-            (2) 只有encoder_output: bsz x max_len x hidden_size, 这种情况下hidden和cell使用0初始化
-        :param torch.ByteTensor encoder_mask: bsz x max_len, 为0的位置是padding, 用来指示source中哪些不需要attend
+        :param encoder_output: ``encoder`` 的输出，可以有两种情况：
+                
+                - 输入一个 :class:`tuple`，包含三个内容 ``(encoder_output, (hidden, cell))``，其中 ``encoder_output`` 形状为
+                  ``[batch_size, max_len, hidden_size]``， ``hidden`` 形状为 ``[batch_size, hidden_size]``， ``cell`` 形状为
+                  ``[batch_size, hidden_size]`` ，一般使用 :class:`~fastNLP.modules.torch.encoder.LSTMSeq2SeqEncoder` 最后一层的
+                  ``hidden state`` 和 ``cell state`` 来赋值这两个值。
+                - 只有形状为 ``[batch_size, max_len, hidden_size]`` 的 ``encoder_output``, 这种情况下 ``hidden`` 和 ``cell``
+                  使用 **0** 初始化。
+        :param encoder_mask: ``[batch_size, max_len]]``，为 **0** 的位置是 padding, 用来指示输入中哪些不需要 attend 。
         :return:
         """
         if not isinstance(encoder_output, torch.Tensor):
@@ -233,14 +241,15 @@ class LSTMSeq2SeqDecoder(Seq2SeqDecoder):
 
 class TransformerSeq2SeqDecoderLayer(nn.Module):
     """
+    **Transformer** 的 Decoder 层
 
-    :param int d_model: 输入、输出的维度
-    :param int n_head: 多少个head，需要能被d_model整除
-    :param int dim_ff:
-    :param float dropout:
-    :param int layer_idx: layer的编号
+    :param d_model: 输入、输出的维度
+    :param n_head: **多头注意力** head 的数目，需要能被 ``d_model`` 整除
+    :param dim_ff:  FFN 中间映射的维度
+    :param dropout: Dropout 的大小
+    :param layer_idx: layer的编号
     """
-    def __init__(self, d_model = 512, n_head = 8, dim_ff = 2048, dropout = 0.1, layer_idx = None):
+    def __init__(self, d_model: int = 512, n_head: int = 8, dim_ff: int = 2048, dropout: float = 0.1, layer_idx: int = None):
         super().__init__()
         self.d_model = d_model
         self.n_head = n_head
@@ -262,14 +271,14 @@ class TransformerSeq2SeqDecoderLayer(nn.Module):
 
         self.final_layer_norm = nn.LayerNorm(self.d_model)
 
-    def forward(self, x, encoder_output, encoder_mask=None, self_attn_mask=None, state=None):
+    def forward(self, x, encoder_output, encoder_mask=None, self_attn_mask=None, state: TransformerState=None):
         """
 
-        :param x: (batch, seq_len, dim), decoder端的输入
-        :param encoder_output: (batch,src_seq_len,dim), encoder的输出
-        :param encoder_mask: batch,src_seq_len, 为1的地方需要attend
-        :param self_attn_mask: seq_len, seq_len，下三角的mask矩阵，只在训练时传入
-        :param TransformerState state: 只在inference阶段传入
+        :param x: ``decoder`` 端的输入，形状为 ``[batch_size, seq_len, dim]`` 
+        :param encoder_output: ``encoder`` 的输出，形状为 ``[batch_size, src_seq_len, dim]``
+        :param encoder_mask: 掩码，形状为 ``[batch_size, src_seq_len]``，为 **1** 的地方表示需要 attend
+        :param self_attn_mask: 下三角的mask矩阵，只在训练时传入。形状为 ``[seq_len, seq_len]``
+        :param state: 只在 inference 阶段传入，记录了 ``encoder`` 和 ``decoder`` 的状态
         :return:
         """
 
@@ -307,16 +316,23 @@ class TransformerSeq2SeqDecoderLayer(nn.Module):
 
 class TransformerSeq2SeqDecoder(Seq2SeqDecoder):
     """
-
-    :param embed: 输入token的embedding
-    :param nn.Module pos_embed: 位置embedding
-    :param int d_model: 输出、输出的大小
-    :param int num_layers: 多少层
-    :param int n_head: 多少个head
-    :param int dim_ff: FFN 的中间大小
-    :param float dropout: Self-Attention和FFN中的dropout的大小
-    :param bool bind_decoder_input_output_embed: 是否将输出层和输入层的词向量绑定在一起（即为同一个），若embed为StaticEmbedding，
-        则StaticEmbedding的vocab不能包含no_create_entry的token，同时StaticEmbedding初始化时lower为False, min_freq=1.
+    **Transformer** 的 Decoder
+
+    :param embed: ``decoder`` 输入的 embedding，支持以下几种输入类型：
+
+            - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+            - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+            - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+            - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+    :param pos_embed: 位置 embedding
+    :param d_model: 输入、输出的维度
+    :param num_layers: :class:`TransformerSeq2SeqDecoderLayer` 的层数
+    :param n_head: **多头注意力** head 的数目，需要能被 ``d_model`` 整除
+    :param dim_ff: FFN 中间映射的维度
+    :param dropout: :class:`~fastNLP.modules.torch.decoder.SelfAttention` 和 FFN 中的 dropout 的大小
+    :param bind_decoder_input_output_embed: ``decoder`` 的输出 embedding 是否与其输入 embedding 是一样的权重（即为同一个），若 ``embed`` 为 
+        :class:`~fastNLP.embeddings.StaticEmbedding`，则 ``StaticEmbedding`` 的 ``vocab`` 不能包含 ``no_create_entry`` 的 token ，同时
+        ``StaticEmbedding`` 初始化时 ``lower`` 为 ``False``，``min_freq=1``。
     """
     def __init__(self, embed: Union[nn.Module, StaticEmbedding, Tuple[int, int]], pos_embed: nn.Module = None,
                  d_model = 512, num_layers=6, n_head = 8, dim_ff = 2048, dropout = 0.1,
@@ -346,13 +362,14 @@ class TransformerSeq2SeqDecoder(Seq2SeqDecoder):
         self.layer_norm = nn.LayerNorm(d_model)
         self.output_fc = nn.Linear(self.d_model, self.embed.embedding_dim)
 
-    def forward(self, tokens, state, return_attention=False):
+    def forward(self, tokens: torch.LongTensor, state: TransformerState, return_attention: bool=False):
         """
 
-        :param torch.LongTensor tokens: batch x tgt_len，decode的词
-        :param TransformerState state: 用于记录encoder的输出以及decode状态的对象，可以通过init_state()获取
-        :param bool return_attention: 是否返回对encoder结果的attention score
-        :return: bsz x max_len x vocab_size; 如果return_attention=True, 还会返回bsz x max_len x encode_length
+        :param tokens: 用于解码的词，形状为 ``[batch_size, tgt_len]``
+        :param state: 用于记录 ``encoder`` 的输出以及 ``decode`` 状态的对象，可以通过 :meth:`init_state` 获取
+        :param return_attention: 是否返回对 ``encoder`` 结果的 attention score
+        :return: 形状为 ``[batch_size, max_len, vocab_size]`` 的结果。如果 ``return_attention=True`` 则返回一个元组，一个元素为结果，第二个结果为
+            注意力权重，形状为 ``[batch_size, max_len, encode_length]``
         """
 
         encoder_output = state.encoder_output
@@ -391,13 +408,13 @@ class TransformerSeq2SeqDecoder(Seq2SeqDecoder):
             return feats, attn_weight
         return feats
 
-    def init_state(self, encoder_output, encoder_mask):
+    def init_state(self, encoder_output: torch.FloatTensor, encoder_mask: torch.ByteTensor) -> TransformerState:
         """
-        初始化一个TransformerState用于forward
+        初始化一个 :class:`~fastNLP.modules.torch.decoder.TransformerState`` 用于 :meth:`forward`
 
-        :param torch.FloatTensor encoder_output: bsz x max_len x d_model, encoder的输出
-        :param torch.ByteTensor encoder_mask: bsz x max_len, 为1的位置需要attend。
-        :return: TransformerState
+        :param encoder_output: ``encoder`` 的输出，形状为 ``[batch_size, max_len, d_model]``
+        :param encoder_mask: ``[batch_size, max_len]]``，为 **0** 的位置是 padding, 用来指示输入中哪些不需要 attend 。
+        :return:
         """
         if isinstance(encoder_output, torch.Tensor):
             encoder_output = encoder_output

fastNLP/modules/torch/decoder/seq2seq_state.py

@@ -9,21 +9,22 @@ __all__ = [
     "TransformerState"
 ]
 
-from typing import Union
+from typing import Union, List, Tuple
 import torch
 
 
 class State:
-    def __init__(self, encoder_output=None, encoder_mask=None, **kwargs):
-        """
-        每个Decoder都有对应的State对象用来承载encoder的输出以及当前时刻之前的decode状态。
-
-        :param Union[torch.Tensor, list, tuple] encoder_output: 如果不为None，内部元素需要为torch.Tensor, 默认其中第一维是batch
-            维度
-        :param Union[torch.Tensor, list, tuple] encoder_mask: 如果部位None，内部元素需要torch.Tensor, 默认其中第一维是batch
-            维度
-        :param kwargs:
-        """
+    """
+    每个 ``Decoder`` 都有对应的 :class:`State` 对象用来承载 ``encoder`` 的输出以及当前时刻之前的 ``decode`` 状态。
+
+    :param encoder_output: 如果不为 ``None`` ，内部元素需要为 :class:`torch.Tensor`，默认其中第一维是 ``batch_size``
+        维度
+    :param encoder_mask: 如果部位 ``None``，内部元素需要为 :class:`torch.Tensor`，默认其中第一维是 ``batch_size``
+        维度
+    :param kwargs:
+    """
+    def __init__(self, encoder_output: Union[torch.Tensor, List, Tuple]=None, 
+                encoder_mask: Union[torch.Tensor, List, Tuple]=None, **kwargs):
         self.encoder_output = encoder_output
         self.encoder_mask = encoder_mask
         self._decode_length = 0
@@ -31,9 +32,7 @@ class State:
     @property
     def num_samples(self):
         """
-        返回的State中包含的是多少个sample的encoder状态，主要用于Generate的时候确定batch的大小。
-
-        :return:
+        返回的 State 中包含的是多少个 sample 的 encoder 状态，主要用于 Generate 的时候确定 batch_size 的大小。
         """
         if self.encoder_output is not None:
             return self.encoder_output.size(0)
@@ -43,9 +42,7 @@ class State:
     @property
     def decode_length(self):
         """
-        当前Decode到哪个token了，decoder只会从decode_length之后的token开始decode, 为0说明还没开始decode。
-
-        :return:
+        当前 Decode 到哪个 token 了，decoder 只会从 decode_length 之后的 token 开始 decode, 为 **0** 说明还没开始 decode。
         """
         return self._decode_length
 
@@ -79,26 +76,27 @@ class State:
 
 
 class LSTMState(State):
-    def __init__(self, encoder_output, encoder_mask, hidden, cell):
-        """
-        LSTMDecoder对应的State，保存encoder的输出以及LSTM解码过程中的一些中间状态
-
-        :param torch.FloatTensor encoder_output: bsz x src_seq_len x encode_output_size，encoder的输出
-        :param torch.BoolTensor encoder_mask: bsz x src_seq_len, 为0的地方是padding
-        :param torch.FloatTensor hidden: num_layers x bsz x hidden_size, 上个时刻的hidden状态
-        :param torch.FloatTensor cell: num_layers x bsz x hidden_size, 上个时刻的cell状态
-        """
+    """
+    :class:`~fastNLP.modules.torch.decoder.LSTMSeq2SeqDecoder` 对应的 :class:`State`，保存 ``encoder`` 的输出以及 ``LSTM`` 解码过程中的一些中间状态
+
+    :param encoder_output: ``encoder`` 的输出，形状为 ``[batch_size, src_seq_len, encode_output_size]``
+    :param encoder_mask: 掩码，形状为 ``[batch_size, src_seq_len]``，为 **1** 的地方表示需要 attend
+    :param hidden: 上个时刻的 ``hidden`` 状态，形状为 ``[num_layers, batch_size, hidden_size]``
+    :param cell: 上个时刻的 ``cell`` 状态，形状为 ``[num_layers, batch_size, hidden_size]``
+    """
+    def __init__(self, encoder_output: torch.FloatTensor, encoder_mask: torch.BoolTensor, hidden: torch.FloatTensor, cell: torch.FloatTensor):
         super().__init__(encoder_output, encoder_mask)
         self.hidden = hidden
         self.cell = cell
         self._input_feed = hidden[0]  # 默认是上一个时刻的输出
 
     @property
-    def input_feed(self):
+    def input_feed(self) -> torch.FloatTensor:
         """
-        LSTMDecoder中每个时刻的输入会把上个token的embedding和input_feed拼接起来输入到下个时刻，在LSTMDecoder不使用attention时，
-            input_feed即上个时刻的hidden state, 否则是attention layer的输出。
-        :return: torch.FloatTensor, bsz x hidden_size
+        :class:`~fastNLP.modules.torch.decoder.LSTMSeq2SeqDecoder` 中每个时刻的输入会把上个 token 的 embedding 和 ``input_feed`` 拼接起来输入到下个时刻，在
+        :class:`~fastNLP.modules.torch.decoder.LSTMSeq2SeqDecoder` 不使用 ``attention`` 时，``input_feed`` 即上个时刻的 ``hidden state``，否则是 ``attention layer`` 的输出。
+        
+        :return: ``[batch_size, hidden_size]``
         """
         return self._input_feed
 
@@ -115,14 +113,14 @@ class LSTMState(State):
 
 
 class TransformerState(State):
-    def __init__(self, encoder_output, encoder_mask, num_decoder_layer):
-        """
-        与TransformerSeq2SeqDecoder对应的State，
-
-        :param torch.FloatTensor encoder_output: bsz x encode_max_len x encoder_output_size, encoder的输出
-        :param torch.ByteTensor encoder_mask: bsz x encode_max_len 为1的地方需要attend
-        :param int num_decoder_layer: decode有多少层
-        """
+    """
+    与 :class:`~fastNLP.modules.torch.decoder.TransformerSeq2SeqDecoder` 对应的 :class:`State`。
+
+    :param encoder_output: ``encoder`` 的输出，形状为 ``[batch_size, encode_max_len, encode_output_size]``，
+    :param encoder_mask: 掩码，形状为 ``[batch_size, encode_max_len]``，为 **1** 的地方表示需要 attend
+    :param num_decoder_layer: decoder 层数
+    """
+    def __init__(self, encoder_output: torch.FloatTensor, encoder_mask: torch.FloatTensor, num_decoder_layer: int):
         super().__init__(encoder_output, encoder_mask)
         self.encoder_key = [None] * num_decoder_layer  # 每一个元素 bsz x encoder_max_len x key_dim
         self.encoder_value = [None] * num_decoder_layer  # 每一个元素 bsz x encoder_max_len x value_dim

fastNLP/modules/torch/dropout.py

@@ -1,5 +1,3 @@
-r"""undocumented"""
-
 __all__ = [
     "TimestepDropout"
 ]
@@ -9,8 +7,8 @@ import torch
 
 class TimestepDropout(torch.nn.Dropout):
     r"""
-    传入参数的shape为 ``(batch_size, num_timesteps, embedding_dim)``
-    使用同一个shape为 ``(batch_size, embedding_dim)`` 的mask在每个timestamp上做dropout。
+    传入参数的 shape 为 ``(batch_size, num_timesteps, embedding_dim)``
+    使用同一个 shape 为 ``(batch_size, embedding_dim)`` 的 mask 在每个 timestamp 上做 dropout。
     """
     
     def forward(self, x):