diff --git a/fastNLP/core/controllers/evaluator.py b/fastNLP/core/controllers/evaluator.py index 7ef15392..faf5d1aa 100644 --- a/fastNLP/core/controllers/evaluator.py +++ b/fastNLP/core/controllers/evaluator.py @@ -77,9 +77,17 @@ class Evaluator: :param evaluate_batch_step_fn: 定制每次 evaluate batch 执行的函数。该函数应接受的两个参数为 `evaluator` 和 `batch`， 不需要有返回值；可以参考 :meth:`~fastNLP.core.controllers.loops.EvaluateBatchLoop.batch_step_fn` 函数。 - :param evaluate_fn: 用来控制 `Evaluator` 在评测的前向传播过程中是调用哪一个函数，例如是 `model.evaluate_step` 还是 - `model.forward`；(1) 如果该值是 None，那么我们会默认使用 `evaluate_step` 函数，如果在模型中没有 - 找到该方法，则使用 `model.forward` 函数；(2) 如果为 str 类型，则尝试从 model 中寻找该方法，找不到则报错。 + :param evaluate_fn: 用来控制 ``Evaluator`` 在评测时调用模型的哪一个函数，例如是 ``evaluate_step`` 还是框架默认的前向接口； + 默认为 ``None``，如果该值是 ``None``，那么我们会默认使用 ``evaluate_step`` , 如果在模型的定义类中没有找到该方法， + 则使用模型默认的前向传播函数，例如对于 pytorch 来说就是 ``forward``。 + + .. note:: + 查找逻辑如下所示： + + 1. 如果 ``evaluate_fn`` 为 None，那么在 model 的类 Model 中寻找方法 ``Model.evaluate_step``;如果没有找到，那么默认使用 ``Model.forward``； + 2. 如果 ``evaluate_fn`` 为一个字符串，例如 'my_step_fn'，那么我们首先会在 model 的类 Model 中寻找方法 ``Model.my_step_fn``， + 如果没有找到，那么会直接报错； + :param input_mapping: 应当为一个字典或者一个函数，表示在当前 step 拿到一个 batch 的数据后，应当做怎样的映射处理： 1. 如果 ``input_mapping`` 是一个字典: diff --git a/fastNLP/core/controllers/trainer.py b/fastNLP/core/controllers/trainer.py index 1b650a7d..73b6d34b 100644 --- a/fastNLP/core/controllers/trainer.py +++ b/fastNLP/core/controllers/trainer.py @@ -38,306 +38,317 @@ from fastNLP.core.utils.exceptions import EarlyStopException class Trainer(TrainerEventTrigger): - _custom_callbacks: dict = defaultdict(list) + r""" + 用于支持快速训练的训练器。 - def __init__( - self, - model, - driver, - train_dataloader, - optimizers, - device: Optional[Union[int, List[int], str]] = "cpu", - n_epochs: int = 20, - evaluate_dataloaders=None, - batch_step_fn: Optional[Callable] = None, - evaluate_batch_step_fn: Optional[Callable] = None, - train_fn: Optional[str] = None, - evaluate_fn: Optional[str] = None, - callbacks: Union[List[Callback], Callback, None] = None, - metrics: Optional[dict] = None, - evaluate_every: Optional[Union[int, Callable]] = -1, - input_mapping: Optional[Union[Callable, Dict]] = None, - output_mapping: Optional[Union[Callable, Dict]] = None, - model_wo_auto_param_call: bool = False, - accumulation_steps: int = 1, - fp16: bool = False, - monitor: Union[str, Callable] = None, - larger_better: bool = True, - marker: Optional[str] = None, - **kwargs - ): - r""" - :param model: 训练所需要的模型，例如 ``torch.nn.Module``； + :param model: 训练所需要的模型，例如 ``torch.nn.Module``； + + .. note:: + + 当使用 pytorch 时，注意参数 ``model`` 在大多数情况下为 ``nn.Module``。但是您仍能够通过使用一些特定的组合来使用情况，如下所示： + + 1. 当希望使用 ``DataParallel`` 时，您应当使用 ``TorchSingleDriver``，意味着您在初始化 ``Trainer`` 时参数 ``device`` 不应当为 + 一个 ``List``； + + 2. 当您选择自己初始化 ``init_process_group`` 时（这种情况要求您传入的 ``model`` 参数一定为 ``DistributedDataParallel``）， + 您应当使用 ``TorchDDPDriver``，意味着您需要通过 ``python -m torch.distributed.launch`` 的方式来启动训练，此时参数 ``device`` + 应当设置为 None（此时我们会忽略该参数），具体见下面对于参数 ``device`` 的更详细的解释。 + + :param driver: 训练模型所使用的具体的驱动模式，应当为以下选择中的一个：["torch"]，之后我们会加入 jittor、paddle 等 + 国产框架的训练模式；其中 "torch" 表示使用 ``TorchSingleDriver`` 或者 ``TorchDDPDriver``，具体使用哪一种取决于参数 ``device`` + 的设置； + :param train_dataloader: 训练数据集，注意其必须是单独的一个数据集，不能是 List 或者 Dict； + :param optimizers: 训练所需要的优化器；可以是单独的一个优化器实例，也可以是多个优化器组成的 List； + :param device: 该参数用来指定具体训练时使用的机器；注意当该参数仅当您通过 `torch.distributed.launch/run` 启动时可以为 None， + 此时 fastNLP 不会对模型和数据进行设备之间的移动处理，但是你可以通过参数 `input_mapping` 和 `output_mapping` 来实现设备之间 + 数据迁移的工作（通过这两个参数传入两个处理数据的函数）；同时你也可以通过在 kwargs 添加参数 "data_device" 来让我们帮助您将数据 + 迁移到指定的机器上（注意这种情况理应只出现在用户在 Trainer 实例化前自己构造 DDP 的场景）； + + device 的可选输入如下所示： + + * *str*: 例如 'cpu', 'cuda', 'cuda:0', 'cuda:1' 等； + * *torch.device*: 例如 'torch.device("cuda:0")'； + * *int*: 将使用 ``device_id`` 为该值的 ``gpu`` 进行训练；如果值为 -1，那么默认使用全部的显卡，此时使用的 driver 实例是 `TorchDDPDriver`； + * *list(int)*: 如果多于 1 个device，应当通过该种方式进行设定；注意此时我们一定会使用 ``TorchDDPDriver``，不管您传入的列表的长度是 1 还是其它值； + * *None*: 仅当用户自己通过训练框架提供的并行训练启动脚本开启 ddp 进程时为 None； .. note:: - 当使用 pytorch 时，注意参数 ``model`` 在大多数情况下为 ``nn.Module``。但是您仍能够通过使用一些特定的组合来使用情况，如下所示： + 如果希望使用 ``TorchDDPDriver``，在初始化 ``Trainer`` 时您应当使用:: - 1. 当希望使用 ``DataParallel`` 时，您应当使用 ``TorchSingleDriver``，意味着您在初始化 ``Trainer`` 时参数 ``device`` 不应当为 - 一个 ``List``； + Trainer(driver="torch", device=[0, 1]) - 2. 当您选择自己初始化 ``init_process_group`` 时（这种情况要求您传入的 ``model`` 参数一定为 ``DistributedDataParallel``）， - 您应当使用 ``TorchDDPDriver``，意味着您需要通过 ``python -m torch.distributed.launch`` 的方式来启动训练，此时参数 ``device`` - 应当设置为 None（此时我们会忽略该参数），具体见下面对于参数 ``device`` 的更详细的解释。 + 注意如果这时 ``device=[0]``，我们仍旧会使用 ``TorchDDPDriver``。 - :param driver: 训练模型所使用的具体的驱动模式，应当为以下选择中的一个：``["torch", "jittor", "paddle"]`` - 其中 "torch" 表示使用 ``TorchSingleDriver`` 或者 ``TorchDDPDriver``，具体使用哪一种取决于参数 ``device`` - 的设置。 - :param train_dataloader: 训练数据集，注意其必须是单独的一个数据集，不能是 List 或者 Dict； - :param optimizers: 训练所需要的优化器；可以是单独的一个优化器实例，也可以是多个优化器组成的 List； - :param device: 该参数用来指定具体训练时使用的机器；注意当该参数仅当您通过 `torch.distributed.launch/run` 启动时可以为 None， - 此时 fastNLP 不会对模型和数据进行设备之间的移动处理，但是你可以通过参数 `input_mapping` 和 `output_mapping` 来实现设备之间 - 数据迁移的工作（通过这两个参数传入两个处理数据的函数）；同时你也可以通过在 kwargs 添加参数 "data_device" 来让我们帮助您将数据 - 迁移到指定的机器上（注意这种情况理应只出现在用户在 Trainer 实例化前自己构造 DDP 的场景）； + 如果希望使用 ``TorchSingleDriver``，则在初始化 ``Trainer`` 时您应当使用:: - device 的可选输入如下所示： + Trainer(driver="torch", device=0) - * *str*: 例如 'cpu', 'cuda', 'cuda:0', 'cuda:1' 等； - * *torch.device*: 例如 'torch.device("cuda:0")'； - * *int*: 将使用 ``device_id`` 为该值的 ``gpu`` 进行训练；如果值为 -1，那么默认使用全部的显卡，此时使用的 driver 实例是 `TorchDDPDriver`； - * *list(int)*: 如果多于 1 个device，应当通过该种方式进行设定；注意此时我们一定会使用 ``TorchDDPDriver``，不管您传入的列表的长度是 1 还是其它值； - * *None*: 仅当用户自己通过训练框架提供的并行训练启动脚本开启 ddp 进程时为 None； + .. warning:: - .. note:: + 注意参数 ``device`` 仅当您通过 pytorch 或者其它训练框架自身的并行训练启动脚本启动 ddp 训练时才允许为 ``None``！ - 如果希望使用 ``TorchDDPDriver``，在初始化 ``Trainer`` 时您应当使用:: + 例如，当您使用:: - Trainer(driver="torch", device=[0, 1]) + python -m torch.distributed.launch --nproc_per_node 2 train.py - 注意如果这时 ``device=[0]``，我们仍旧会使用 ``TorchDDPDriver``。 + 来使用 ``TorchDDPDriver`` 时，此时参数 ``device`` 不再有效（不管您是否自己初始化 ``init_process_group``），我们将直接 + 通过 ``torch.device(f"cuda:{local_rank}")`` 来获取当前进程所使用的的具体的 gpu 设备。因此此时您需要使用 ``os.environ["CUDA_VISIBLE_DEVICES"]`` + 来指定要使用的具体的 gpu 设备。 - 如果希望使用 ``TorchSingleDriver``，则在初始化 ``Trainer`` 时您应当使用:: + 另一点需要注意的是，当您没有选择自己初始化 ``init_process_group`` 时，我们仍旧会帮助您把模型和数据迁移到当前进程所使用的 + 具体的 gpu 设备上。但是如果您选择自己在 ``Trainer`` 初始化前（意味着在 ``driver`` 的 ``setup`` 前）初始化 ``init_process_group``， + 那么对于模型的迁移应当完全由您自己来完成。此时对于数据的迁移，如果您在 ``Trainer`` 初始化时指定了参数 ``data_device``，那么 + 我们会将数据迁移到 ``data_device`` 上；如果其为 None，那么将数据迁移到正确的设备上应当由您自己来完成。 - Trainer(driver="torch", device=0) + 对于使用 ``TorchDDPDriver`` 的更多细节，请见 :class:`fastNLP.core.drivers.torch_driver.TorchDDPDriver`。 - .. warning:: + :param n_epochs: 训练总共的 epoch 的数量，默认为 20； + :param evaluate_dataloaders: 验证数据集，其可以是单独的一个数据集，也可以是多个数据集；当为多个数据集时，注意其必须是 Dict；默认 + 为 None； + :param batch_step_fn: 定制每次训练时前向运行一个 batch 的数据所执行的函数。该函数应接受两个参数为 ``trainer`` 和 ``batch``， + 不需要要返回值；更详细的使用位置和说明请见 :meth:`fastNLP.core.controllers.TrainBatchLoop.batch_step_fn`； + :param evaluate_batch_step_fn: 定制每次验证时前向运行一个 batch 的数据所执行的函数。该函数应接受的两个参数为 ``evaluator`` 和 ``batch``， + 不需要有返回值；可以参考 :meth:`fastNLP.core.controllers.EvaluateBatchLoop.batch_step_fn`； + :param train_fn: 用来控制 ``Trainer`` 在训练的前向传播过程中是调用模型的哪一个函数，例如是 ``train_step`` 还是框架默认的前向接口； + 默认为 ``None``，如果该值是 ``None``，那么我们会默认使用 ``train_step`` 当做前向传播的函数，如果在模型的定义类中没有找到该方法， + 则使用模型默认的前向传播函数，例如对于 pytorch 来说就是 ``forward``。 - 注意参数 ``device`` 仅当您通过 pytorch 或者其它训练框架自身的并行训练启动脚本启动 ddp 训练时才允许为 ``None``！ + .. note:: + 在 fastNLP 中，对于训练时使用的前向传播函数的查找逻辑如下所示： - 例如，当您使用:: + 1. 如果 ``train_fn`` 为 None，那么在 model 的类 Model 中寻找方法 ``Model.train_step``;如果没有找到，那么默认使用 ``Model.forward``； + 2. 如果 ``train_fn`` 为一个字符串，例如 'my_step_fn'，那么我们首先会在 model 的类 Model 中寻找方法 ``Model.my_step_fn``， + 如果没有找到，那么会直接报错； - python -m torch.distributed.launch --nproc_per_node 2 train.py + :param evaluate_fn: 用来控制 ``Trainer`` 中内置的 ``Evaluator`` 在验证的前向传播过程中是调用模型的哪一个函数，应当为 ``None`` + 或者一个字符串；其使用方式和 train_fn 类似；具体可见 :class:`fastNLP.core.controllers.Evaluator`； + :param callbacks: 训练当中触发的 callback 类，该参数应当为一个列表，其中的每一个元素都应当继承 ``Callback`` 类；具体可见 + :class:`fastNLP.core.callbacks.Callback`； + :param metrics: 用于传给 ``Trainer`` 内部的 ``Evaluator`` 实例来进行训练过程中的验证。其应当为一个字典，其中 key 表示 monitor， + 例如 {"acc1": AccMetric(), "acc2": AccMetric()}； - 来使用 ``TorchDDPDriver`` 时，此时参数 ``device`` 不再有效（不管您是否自己初始化 ``init_process_group``），我们将直接 - 通过 ``torch.device(f"cuda:{local_rank}")`` 来获取当前进程所使用的的具体的 gpu 设备。因此此时您需要使用 ``os.environ["CUDA_VISIBLE_DEVICES"]`` - 来指定要使用的具体的 gpu 设备。 + 目前我们支持的 ``metric`` 的种类有以下几种： - 另一点需要注意的是，当您没有选择自己初始化 ``init_process_group`` 时，我们仍旧会帮助您把模型和数据迁移到当前进程所使用的 - 具体的 gpu 设备上。但是如果您选择自己在 ``Trainer`` 初始化前（意味着在 ``driver`` 的 ``setup`` 前）初始化 ``init_process_group``， - 那么对于模型的迁移应当完全由您自己来完成。此时对于数据的迁移，如果您在 ``Trainer`` 初始化时指定了参数 ``data_device``，那么 - 我们会将数据迁移到 ``data_device`` 上；如果其为 None，那么将数据迁移到正确的设备上应当由您自己来完成。 + 1. fastNLP 自己的 ``metric``：详见 :class:`fastNLP.core.metrics.Metric`； + 2. torchmetrics； + 3. allennlp.training.metrics； + 4. paddle.metric； - 对于使用 ``TorchDDPDriver`` 的更多细节，请见 :class:`fastNLP.core.drivers.torch_driver.TorchDDPDriver`。 + :param evaluate_every: 用来控制 ``Trainer`` 内部的 ``Evaluator`` 验证的频率，其可以为负数、正数或者函数： - :param n_epochs: 训练总共的 epoch 的数量，默认为 20； - :param evaluate_dataloaders: 验证数据集，其可以是单独的一个数据集，也可以是多个数据集；当为多个数据集时，注意其必须是 Dict；默认 - 为 None； - :param batch_step_fn: 定制每次训练时前向运行一个 batch 的数据所执行的函数。该函数应接受两个参数为 ``trainer`` 和 ``batch``， - 不需要要返回值；更详细的使用位置和说明请见 :meth:`fastNLP.core.controllers.TrainBatchLoop.batch_step_fn`； - :param evaluate_batch_step_fn: 定制每次验证时前向运行一个 batch 的数据所执行的函数。该函数应接受的两个参数为 ``evaluator`` 和 ``batch``， - 不需要有返回值；可以参考 :meth:`fastNLP.core.controllers.TrainBatchLoop.batch_step_fn`； - :param train_fn: 用来控制 ``Trainer`` 在训练的前向传播过程中是调用模型的哪一个函数，例如是 ``train_step`` 还是 ``forward``； - 默认为 ``None``，如果该值是 ``None``，那么我们会默认使用 ``train_step`` 当做前向传播的函数，如果在模型的定义类中没有找到该方法， - 则使用模型默认的前向传播函数，例如对于 pytorch 来说就是 ``forward``。 + 1. 为负数时表示每隔几个 ``epoch`` evaluate 一次； + 2. 为正数则表示每隔几个 ``batch`` evaluate 一次； + 3. 为函数时表示用户自己传入的用于控制 evaluate 的频率的函数，该函数的应该接受当前 trainer 对象作为参数，并 + 返回一个 bool 值，返回为 True 说明需要进行 evaluate ；将在每个 ``batch`` 结束后调用该函数判断是否需要 evaluate； - .. note:: - 在 fastNLP 中，对于训练时使用的前向传播函数的查找逻辑如下所示： + .. note:: - 1. 如果 ``train_fn`` 为 None，那么在 model 的类 Model 中寻找方法 ``Model.train_step``;如果没有找到，那么默认使用 ``Model.forward``； - 2. 如果 ``train_fn`` 为一个字符串，例如 'my_step_fn'，那么我们首先会在 model 的类 Model 中寻找方法 ``Model.my_step_fn``， - 如果没有找到，那么会直接报错； + 如果参数 ``evaluate_every`` 为函数，其应当类似： - :param evaluate_fn: 用来控制 ``Trainer`` 中内置的 ``Evaluator`` 在验证的前向传播过程中是调用模型的哪一个函数，应当为 ``None`` - 或者一个字符串；其使用方式和 train_fn 类似；具体可见 :class:`fastNLP.core.controllers.Evaluator`； - :param callbacks: 训练当中触发的 callback 类，该参数应当为一个列表，其中的每一个元素都应当继承 ``Callback`` 类；具体可见 - :class:`fastNLP.core.callbacks.Callback`； - :param metrics: 用于传给 ``Trainer`` 内部的 ``Evaluator`` 实例来进行训练过程中的验证。其应当为一个字典，其中 key 表示 monitor， - 例如 {"acc1": AccMetric(), "acc2": AccMetric()}； + >>> def my_evaluate_every(trainer) -> bool: + ... if (trainer.global_forward_batches+1) % 1000 == 0: + ... return True + ... else: + ... return False - 目前我们支持的 ``metric`` 的种类有以下几种： + 该函数表示当每经过 1000 个 batch，``Trainer`` 中内置的 ``Evaluator`` 就会验证一次； - 1. fastNLP 自己的 ``metric``：详见 :class:`fastNLP.core.metrics.Metric`； - 2. torchmetrics； - 3. allennlp.training.metrics； - 4. paddle.metric； + 另一个需要注意的事情在于该函数会在每一次 batch 的结尾进行调用，当该函数返回 ``True`` 时，``Evaluator`` 才会进行验证； - :param evaluate_every: 用来控制 ``Trainer`` 内部的 ``Evaluator`` 验证的频率，其可以为负数、正数或者函数： + :param input_mapping: 应当为一个字典或者一个函数，表示在当前 step 拿到一个 batch 的训练数据后，应当做怎样的映射处理： - 1. 为负数时表示每隔几个 ``epoch`` evaluate 一次； - 2. 为正数则表示每隔几个 ``batch`` evaluate 一次； - 3. 为函数时表示用户自己传入的用于控制 evaluate 的频率的函数，该函数的应该接受当前 trainer 对象作为参数，并 - 返回一个 bool 值，返回为 True 说明需要进行 evaluate ；将在每个 ``batch`` 结束后调用该函数判断是否需要 evaluate； + 1. 如果 ``input_mapping`` 是一个字典: - .. note:: + 1. 如果此时 batch 也是一个 ``Dict``，那么我们会把 batch 中同样在 ``input_mapping`` 中的 key 修改为 ``input_mapping`` 的对应 ``key`` 的 ``value``； + 2. 如果此时 batch 是一个 ``dataclass``，那么我们会先将其转换为一个 ``Dict``，然后再进行上述转换； + 3. 如果此时 batch 此时是其它类型，那么我们将会直接报错； + 2. 如果 ``input_mapping`` 是一个函数，那么对于取出的 batch，我们将不会做任何处理，而是直接将其传入该函数里； - 如果参数 ``evaluate_every`` 为函数，其应当类似： + 注意该参数会被传进 ``Evaluator`` 中；因此你可以通过该参数来实现将训练数据 batch 移到对应机器上的工作（例如当参数 ``device`` 为 ``None`` 时）； + 如果 ``Trainer`` 和 ``Evaluator`` 需要使用不同的 ``input_mapping``, 请使用 ``train_input_mapping`` 与 ``evaluate_input_mapping`` 分别进行设置。 - >>> def my_evaluate_every(trainer) -> bool: - ... if (trainer.global_forward_batches+1) % 1000 == 0: - ... return True - ... else: - ... return False + :param output_mapping: 应当为一个字典或者函数。作用和 ``input_mapping`` 类似，区别在于其用于转换输出： - 该函数表示当每经过 1000 个 batch，``Trainer`` 中内置的 ``Evaluator`` 就会验证一次； + 1. 如果 ``output_mapping`` 是一个 ``Dict``，那么我们需要模型的输出必须是 ``Dict`` 或者 ``dataclass`` 类型： - 另一个需要注意的事情在于该函数会在每一次 batch 的结尾进行调用，当该函数返回 ``True`` 时，``Evaluator`` 才会进行验证； + 1. 如果此时模型的输出是一个 ``Dict``，那么我们会把输出中同样在 ``output_mapping`` 中的 key 修改为 ``output_mapping`` 的对应 key 的 value； + 2. 如果此时模型的输出是一个 ``dataclass``，那么我们会先将其转换为一个 Dict，然后再进行上述转换； + 2. 如果 ``output_mapping`` 是一个函数，那么我们将会直接将模型的输出传给该函数； - :param input_mapping: 应当为一个字典或者一个函数，表示在当前 step 拿到一个 batch 的训练数据后，应当做怎样的映射处理： + 如果 ``Trainer`` 和 ``Evaluator`` 需要使用不同的 ``output_mapping``, 请使用 ``train_output_mapping`` 与 ``evaluate_output_mapping`` 分别进行设置； - 1. 如果 ``input_mapping`` 是一个字典: + .. note:: - 1. 如果此时 batch 也是一个 ``Dict``，那么我们会把 batch 中同样在 ``input_mapping`` 中的 key 修改为 ``input_mapping`` 的对应 ``key`` 的 ``value``； - 2. 如果此时 batch 是一个 ``dataclass``，那么我们会先将其转换为一个 ``Dict``，然后再进行上述转换； - 3. 如果此时 batch 此时是其它类型，那么我们将会直接报错； - 2. 如果 ``input_mapping`` 是一个函数，那么对于取出的 batch，我们将不会做任何处理，而是直接将其传入该函数里； + ``input_mapping`` 和 ``output_mapping`` 与 fastNLP 的一个特殊的概念 **'参数绑定'** 高度相关，它们的存在也是为了 fastNLP + 中的参数匹配能够正确地运行； - 注意该参数会被传进 ``Evaluator`` 中；因此你可以通过该参数来实现将训练数据 batch 移到对应机器上的工作（例如当参数 ``device`` 为 ``None`` 时）； - 如果 ``Trainer`` 和 ``Evaluator`` 需要使用不同的 ``input_mapping``, 请使用 ``train_input_mapping`` 与 ``evaluate_input_mapping`` 分别进行设置。 + .. todo:: + 之后链接上 参数匹配 的文档； - :param output_mapping: 应当为一个字典或者函数。作用和 ``input_mapping`` 类似，区别在于其用于转换输出： + .. warning:: - 1. 如果 ``output_mapping`` 是一个 ``Dict``，那么我们需要模型的输出必须是 ``Dict`` 或者 ``dataclass`` 类型： + 如果 ``Trainer`` 的参数 ``output_mapping`` 不为 ``None``，请保证其返回的一定是一个字典，并且其中含有关键字 **'loss'**； - 1. 如果此时模型的输出是一个 ``Dict``，那么我们会把输出中同样在 ``output_mapping`` 中的 key 修改为 ``output_mapping`` 的对应 key 的 value； - 2. 如果此时模型的输出是一个 ``dataclass``，那么我们会先将其转换为一个 Dict，然后再进行上述转换； - 2. 如果 ``output_mapping`` 是一个函数，那么我们将会直接将模型的输出传给该函数； + :param model_wo_auto_param_call: 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为； - 如果 ``Trainer`` 和 ``Evaluator`` 需要使用不同的 ``output_mapping``, 请使用 ``train_output_mapping`` 与 ``evaluate_output_mapping`` 分别进行设置； + 1. 如果该值为 ``False``，并且当 batch 为字典时，我们会根据**前向函数**所需要的参数从 batch 中提取对应的对象，然后传入到**前向函数**中； + 2. 如果该值为 ``True``，那么我们会将 batch 直接透传给模型； - .. note:: + .. todo:: + 之后链接上 参数匹配 的文档； - ``input_mapping`` 和 ``output_mapping`` 与 fastNLP 的一个特殊的概念 **'参数绑定'** 高度相关，它们的存在也是为了 fastNLP - 中的参数匹配能够正确地运行； + 函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`； - .. todo:: - 之后链接上 参数匹配 的文档； + :param accumulation_steps: 梯度累积的步数，表示每隔几个 batch 才让优化器迭代一次，默认为 1； + :param fp16: 是否开启混合精度训练，默认为 False； + :param monitor: 对于一些特殊的 ``Callback``，例如 :class:`fastNLP.core.callbacks.CheckpointCallback`，它们需要参数 ``monitor`` + 来从 ``Evaluator`` 的验证结果中获取当前评测的值，从而来判断是否执行一些特殊的操作。例如，对于 ``CheckpointCallback`` 而言，如果我们 + 想要每隔一个 epoch 让 ``Evaluator`` 进行一次验证，然后保存训练以来的最好的结果；那么我们需要这样设置： - .. warning:: + .. code-block:: + + trainer = Trainer( + ..., + metrics={'acc': accMetric()}, + callbacks=[CheckpointCallback( + ..., + monitor='acc', + topk=1 + )] + ) - 如果 ``Trainer`` 的参数 ``output_mapping`` 不为 ``None``，请保证其返回的一定是一个字典，并且其中含有关键字 **'loss'**； + 这意味着对于 ``CheckpointCallback`` 来说，*'acc'* 就是一个监测的指标，用于在 ``Evaluator`` 验证后取出其需要监测的那个指标的值。 - :param model_wo_auto_param_call: 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为； + ``Trainer`` 中的参数 ``monitor`` 的作用在于为没有设置 ``monitor`` 参数但是需要该参数的 *callback* 实例设置该值。关于 ``monitor`` + 参数更详细的说明，请见 :class:`fastNLP.core.callbacks.CheckpointCallback`； - 1. 如果该值为 ``False``，并且当 batch 为字典时，我们会根据**前向函数**所需要的参数从 batch 中提取对应的对象，然后传入到**前向函数**中； - 2. 如果该值为 ``True``，那么我们会将 batch 直接透传给模型； + 注意该参数仅当 ``Trainer`` 内置的 ``Evaluator`` 不为 None 时且有需要该参数但是没有设置该参数的 *callback* 实例才有效； - .. todo:: - 之后链接上 参数匹配 的文档； + :param larger_better: 对于需要参数 ``monitor`` 的 *callback* 来说，``monitor`` 的值是否是越大越好；类似于 ``monitor``，其作用 + 在于为没有设置 ``larger_better`` 参数但是需要该参数的 *callback* 实例设置该值； + + 注意该参数仅当 ``Trainer`` 内置的 ``Evaluator`` 不为 None 时且有需要该参数但是没有设置该参数的 *callback* 实例才有效； + + :param marker: 用于标记一个 ``Trainer`` 实例，从而在用户调用 ``Trainer.on`` 函数时，标记该函数属于哪一个具体的 ``Trainer`` 实例；默认为 None； + + .. note:: + + marker 的使用场景主要在于如果一个脚本中含有多个 ``Trainer`` 实例，并且含有多个使用 ``Trainer.on`` 修饰的函数时，不同的函数属于 + 不同的 ``Trainer`` 实例； + + 此时，通过将修饰器 ``Trainer.on`` 的参数 ``marker`` 和 ``Trainer`` 的参数 ``marker`` 置为相同，就可以使得该函数只会在这一 + ``Trainer`` 实例中被调用；例如， - 函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`； + .. code-block:: - :param accumulation_steps: 梯度累积的步数，表示每隔几个 batch 才让优化器迭代一次，默认为 1； - :param fp16: 是否开启混合精度训练，默认为 False； - :param monitor: 对于一些特殊的 ``Callback``，例如 :class:`fastNLP.core.callbacks.CheckpointCallback`，它们需要参数 ``monitor`` - 来从 ``Evaluator`` 的验证结果中获取当前评测的值，从而来判断是否执行一些特殊的操作。这里设置了 ``monitor`` 则所有的需要 - ``monitor`` 但是没有自己设置的 ``Callback`` 都会使用这个值 + @Trainer.on(Event.on_train_begin(), marker='trainer1') + def fn(trainer): + ... - * 为 ``None`` - 没有 monitor ，默认。 - * 为 ``str`` - 尝试直接使用该名称从 ``evaluation`` 结果中寻找，如果在 ``evaluation`` 结果中没有找到完全一致的名称，将 - 使用 最长公共字符串算法 从 ``evaluation`` 结果中找到最匹配的那个作为 ``monitor`` 。 - * 为 ``Callable`` - 接受参数为 ``evaluation`` 的结果(字典类型)，返回一个 ``float`` 值作为 ``monitor`` 的结果，如果当前结果中没有相关 - 的 ``monitor`` 值请返回 ``None`` 。 + trainer = Trainer( + ..., + marker='trainer1' + ) - 注意该参数仅当传入了 ``evaluate_dataloaders`` 不为 ``None`` 时且有需要该参数但是没有设置该参数的 *Callback* 实例才有意义； + 另一点需要说明的是，如果一个被 ``Trainer.on`` 修饰的函数，其修饰时没有指明 ``marker``，那么会将该函数传给代码位于其之后的 + 第一个 ``Trainer`` 实例，即使该 ``Trainer`` 实例的 marker 不为 None；这一点详见 :meth:`~fastNLP.core.controllers.Trainer.on` - :param larger_better: 对于需要参数 ``monitor`` 的 *callback* 来说，``monitor`` 的值是否是越大越好；类似于 ``monitor``，其作用 - 在于为没有设置 ``larger_better`` 参数但是需要该参数的 *callback* 实例设置该值； + :kwargs: + * *torch_kwargs* -- 用于在指定 ``driver`` 为 'torch' 时设定具体 driver 实例的一些参数： + * ddp_kwargs -- 用于在使用 ``TorchDDPDriver`` 时指定 ``DistributedDataParallel`` 初始化时的参数；例如传入 + {'find_unused_parameters': True} 来解决有参数不参与前向运算导致的报错等； + * set_grad_to_none -- 是否在训练过程中在每一次 optimizer 更新后将 grad 置为 None； + * torch_non_blocking -- 表示用于 pytorch 的 tensor 的 to 方法的参数 non_blocking； + * *paddle_kwargs* -- 用于在指定 ``driver`` 为 'paddle' 时设定具体 driver 实例的一些参数： - 注意该参数仅当 ``Trainer`` 内置的 ``Evaluator`` 不为 None 时且有需要该参数但是没有设置该参数的 *callback* 实例才有效； + * fleet_kwargs -- 用于在使用 ``PaddleFleetDriver`` 时指定 ``DataParallel`` 和 ``fleet`` 初始化时的参数，包括： - :param marker: 用于标记一个 ``Trainer`` 实例，从而在用户调用 ``Trainer.on`` 函数时，标记该函数属于哪一个具体的 ``Trainer`` 实例；默认为 None； + * is_collective -- 是否使用 paddle 集群式的分布式训练方法，目前仅支持为 ``True`` 的情况； + * role_maker -- 初始化 ``fleet`` 分布式训练 API 时使用的 ``RoleMaker`` + * 其它用于初始化 ``DataParallel`` 的参数； + * *data_device* -- 一个具体的 driver 实例中，有 ``model_device`` 和 ``data_device``，前者表示模型所在的设备，后者表示 + 当 ``model_device`` 为 None 时应当将数据迁移到哪个设备； .. note:: - marker 的使用场景主要在于如果一个脚本中含有多个 ``Trainer`` 实例，并且含有多个使用 ``Trainer.on`` 修饰的函数时，不同的函数属于 - 不同的 ``Trainer`` 实例； - - 此时，通过将修饰器 ``Trainer.on`` 的参数 ``marker`` 和 ``Trainer`` 的参数 ``marker`` 置为相同，就可以使得该函数只会在这一 - ``Trainer`` 实例中被调用；例如， - - .. code-block:: - - @Trainer.on(Event.on_train_begin(), marker='trainer1') - def fn(trainer): - ... - - trainer = Trainer( - ..., - marker='trainer1' - ) - - 另一点需要说明的是，如果一个被 ``Trainer.on`` 修饰的函数，其修饰时没有指明 ``marker``，那么会将该函数传给代码位于其之后的 - 第一个 ``Trainer`` 实例，即使该 ``Trainer`` 实例的 marker 不为 None；这一点详见 :meth:`~fastNLP.core.controllers.Trainer.on` - - :kwargs: - * *torch_kwargs* -- 用于在指定 ``driver`` 为 'torch' 时设定具体 driver 实例的一些参数： - - * ddp_kwargs -- 用于在使用 ``TorchDDPDriver`` 时指定 ``DistributedDataParallel`` 初始化时的参数；例如传入 - {'find_unused_parameters': True} 来解决有参数不参与前向运算导致的报错等； - * set_grad_to_none -- 是否在训练过程中在每一次 optimizer 更新后将 grad 置为 None； - * torch_non_blocking -- 表示用于 pytorch 的 tensor 的 to 方法的参数 non_blocking； - * *paddle_kwargs* -- 用于在指定 ``driver`` 为 'paddle' 时设定具体 driver 实例的一些参数： - - * fleet_kwargs -- 用于在使用 ``PaddleFleetDriver`` 时指定 ``DataParallel`` 和 ``fleet`` 初始化时的参数，包括： - - * is_collective -- 是否使用 paddle 集群式的分布式训练方法，目前仅支持为 True 的情况； - * role_maker -- 初始化 ``fleet`` 分布式训练 API 时使用的 ``RoleMaker`` - * 其它用于初始化 ``DataParallel`` 的参数； - * *data_device* -- 一个具体的 driver 实例中，有 ``model_device`` 和 ``data_device``，前者表示模型所在的设备，后者表示 - 当 ``model_device`` 为 None 时应当将数据迁移到哪个设备； - - .. note:: - - 注意您在绝大部分情况下不会用到该参数！ - - 1. 当 driver 实例的 ``model_device`` 不为 None 时，该参数无效； - 2. 对于 pytorch，仅当用户自己通过 ``python -m torch.distributed.launch`` 并且自己初始化 ``init_process_group`` 时， - driver 实例的 ``model_device`` 才会为 None； - 3. 对于 paddle，该参数无效； - - * *use_dist_sampler* -- 表示是否使用分布式的 ``sampler``。在多卡时，分布式 ``sampler`` 将自动决定每张卡上读取的 sample ，使得一个 epoch - 内所有卡的 sample 加起来为一整个数据集的 sample。默认会根据 driver 是否为分布式进行设置。 - * *evaluate_use_dist_sampler* -- 表示在 ``Evaluator`` 中在使用分布式的时候是否将 dataloader 的 ``sampler`` 替换为分布式的 ``sampler``；默认为 ``True``； - * *output_from_new_proc* -- 应当为一个字符串，表示在多进程的 driver 中其它进程的输出流应当被做如何处理；其值应当为以下之一： - ["all", "ignore", "only_error"]；当该参数的值不是以上值时，该值应当表示一个文件夹的名字，我们会将其他 rank 的输出流重定向到 - log 文件中，然后将 log 文件保存在通过该参数值设定的文件夹中；默认为 "only_error"； - - 注意该参数仅当使用分布式的 ``driver`` 时才有效，例如 ``TorchDDPDriver``； - * *progress_bar* -- 以哪种方式显示 progress ，目前支持[None, 'raw', 'rich', 'auto'] 或者 RichCallback, RawTextCallback对象， - 默认为 auto , auto 表示如果检测到当前 terminal 为交互型则使用 RichCallback，否则使用 RawTextCallback对象。如果 - 需要定制 progress bar 的参数，例如打印频率等，可以传入 RichCallback, RawTextCallback 对象。 - * *train_input_mapping* -- 与 input_mapping 一致，但是只用于 ``Trainer`` 中。与 input_mapping 互斥。 - * *train_output_mapping* -- 与 output_mapping 一致，但是只用于 ``Trainer`` 中。与 output_mapping 互斥。 - * *evaluate_input_mapping* -- 与 input_mapping 一致，但是只用于 ``Evaluator`` 中。与 input_mapping 互斥。 - * *evaluate_output_mapping* -- 与 output_mapping 一致，但是只用于 ``Evaluator`` 中。与 output_mapping 互斥。 + 注意您在绝大部分情况下不会用到该参数！ - .. note:: - ``Trainer`` 是通过在内部直接初始化一个 ``Evaluator`` 来进行验证； - ``Trainer`` 内部的 ``Evaluator`` 默认是 None，如果您需要在训练过程中进行验证，你需要保证这几个参数得到正确的传入： + 1. 当 driver 实例的 ``model_device`` 不为 None 时，该参数无效； + 2. 对于 pytorch，仅当用户自己通过 ``python -m torch.distributed.launch`` 并且自己初始化 ``init_process_group`` 时， + driver 实例的 ``model_device`` 才会为 None； + 3. 对于 paddle，该参数无效； - 必须的参数：1. ``metrics``；2. ``evaluate_dataloaders``； + * *use_dist_sampler* -- 表示是否使用分布式的 ``sampler``。在多卡时，分布式 ``sampler`` 将自动决定每张卡上读取的 sample ，使得一个 epoch + 内所有卡的 sample 加起来为一整个数据集的 sample。默认会根据 driver 是否为分布式进行设置。 + * *evaluate_use_dist_sampler* -- 表示在 ``Evaluator`` 中在使用分布式的时候是否将 dataloader 的 ``sampler`` 替换为分布式的 ``sampler``；默认为 ``True``； + * *output_from_new_proc* -- 应当为一个字符串，表示在多进程的 driver 中其它进程的输出流应当被做如何处理；其值应当为以下之一： + ["all", "ignore", "only_error"]；当该参数的值不是以上值时，该值应当表示一个文件夹的名字，我们会将其他 rank 的输出流重定向到 + log 文件中，然后将 log 文件保存在通过该参数值设定的文件夹中；默认为 "only_error"； - 可选的其它参数：1. ``evaluate_batch_step_fn；2. ``evaluate_fn``；3. ``evaluate_every``；4. ``input_mapping``； - 5. ``output_mapping``； 6. ``model_wo_auto_param_call``；7. ``fp16``；8. ``monitor``；9. ``larger_better``； + 注意该参数仅当使用分布式的 ``driver`` 时才有效，例如 ``TorchDDPDriver``； + * *progress_bar* -- 以哪种方式显示 progress ，目前支持[None, 'raw', 'rich', 'auto'] 或者 RichCallback, RawTextCallback对象， + 默认为 auto , auto 表示如果检测到当前 terminal 为交互型则使用 RichCallback，否则使用 RawTextCallback对象。如果 + 需要定制 progress bar 的参数，例如打印频率等，可以传入 RichCallback, RawTextCallback 对象。 + * *train_input_mapping* -- 与 input_mapping 一致，但是只用于 ``Trainer`` 中。与 input_mapping 互斥。 + * *train_output_mapping* -- 与 output_mapping 一致，但是只用于 ``Trainer`` 中。与 output_mapping 互斥。 + * *evaluate_input_mapping* -- 与 input_mapping 一致，但是只用于 ``Evaluator`` 中。与 input_mapping 互斥。 + * *evaluate_output_mapping* -- 与 output_mapping 一致，但是只用于 ``Evaluator`` 中。与 output_mapping 互斥。 - .. warning:: + .. note:: + ``Trainer`` 是通过在内部直接初始化一个 ``Evaluator`` 来进行验证； + ``Trainer`` 内部的 ``Evaluator`` 默认是 None，如果您需要在训练过程中进行验证，你需要保证这几个参数得到正确的传入： - 如果 ``Trainer`` 中内置的 ``Evaluator`` 实例不为 ``None``，那么需要注意 ``Trainer`` 中的一些参数是与 ``Evaluator`` 一致的，它们分别为： + 必须的参数：1. ``metrics``；2. ``evaluate_dataloaders``； - 1. ``Evaluator`` 在初始化时的 ``driver`` 参数是 ``Trainer`` 中已经实例化过的 driver；这一点使得一些参数对于 ``Trainer`` 内部的 - ``Evaluator`` 没有用处，例如 ``device``，``torch_kwargs``，``data_device`` 和 ``output_from_new_proc`` 等； - 2. ``input_mapping``，``output_mapping``，``model_wo_auto_param_call`` 和 ``fp16`` 是 ``Trainer`` 和其内部默认的 - ``Evaluator`` 是一致的； + 可选的其它参数：1. ``evaluate_batch_step_fn；2. ``evaluate_fn``；3. ``evaluate_every``；4. ``input_mapping``； + 5. ``output_mapping``； 6. ``model_wo_auto_param_call``；7. ``fp16``；8. ``monitor``；9. ``larger_better``； - 当然，对于 ``input_mapping`` 和 ``output_mapping``，您可以通过添加 ``kwargs`` 中的参数 ``evaluate_input_mapping`` 和 - ``evaluate_output_mapping`` 来单独为 ``Evaluator`` 进行更细致的订制。 + .. warning:: - 另一方面，注意一些专门独属于 ``Evaluator`` 的参数仅当 ``Evaluator`` 不为 None 时才会生效。 + 如果 ``Trainer`` 中内置的 ``Evaluator`` 实例不为 ``None``，那么需要注意 ``Trainer`` 中的一些参数是与 ``Evaluator`` 一致的，它们分别为： + + 1. ``Evaluator`` 在初始化时的 ``driver`` 参数是 ``Trainer`` 中已经实例化过的 driver；这一点使得一些参数对于 ``Trainer`` 内部的 + ``Evaluator`` 没有用处，例如 ``device``，``torch_kwargs``，``data_device`` 和 ``output_from_new_proc`` 等； + 2. ``input_mapping``，``output_mapping``，``model_wo_auto_param_call`` 和 ``fp16`` 是 ``Trainer`` 和其内部默认的 + ``Evaluator`` 是一致的； + + 当然，对于 ``input_mapping`` 和 ``output_mapping``，您可以通过添加 ``kwargs`` 中的参数 ``evaluate_input_mapping`` 和 + ``evaluate_output_mapping`` 来单独为 ``Evaluator`` 进行更细致的订制。 + + 另一方面，注意一些专门独属于 ``Evaluator`` 的参数仅当 ``Evaluator`` 不为 None 时才会生效。 + + """ + + _custom_callbacks: dict = defaultdict(list) + + def __init__( + self, + model, + driver, + train_dataloader, + optimizers, + device: Optional[Union[int, List[int], str]] = "cpu", + n_epochs: int = 20, + evaluate_dataloaders=None, + batch_step_fn: Optional[Callable] = None, + evaluate_batch_step_fn: Optional[Callable] = None, + train_fn: Optional[str] = None, + evaluate_fn: Optional[str] = None, + callbacks: Union[List[Callback], Callback, None] = None, + metrics: Optional[dict] = None, + evaluate_every: Optional[Union[int, Callable]] = -1, + input_mapping: Optional[Union[Callable, Dict]] = None, + output_mapping: Optional[Union[Callable, Dict]] = None, + model_wo_auto_param_call: bool = False, + accumulation_steps: int = 1, + fp16: bool = False, + monitor: Union[str, Callable] = None, + larger_better: bool = True, + marker: Optional[str] = None, + **kwargs + ): - """ self.model = model self.marker = marker if isinstance(driver, str): @@ -817,8 +828,10 @@ class Trainer(TrainerEventTrigger): def _fetch_matched_fn_callbacks(self): r""" - 因为对于使用装饰器加入的函数 callback，我们是加在类属性中，因此在初始化一个具体的 trainer 实例后，我们需要从 Trainer 的 - callback 类属性中将属于其的 callback 函数拿到，然后加入到 callback_manager 中； + 因为对于使用装饰器加入的函数 callback，我们是加在类属性 ``_custom_callbacks`` 中，因此在初始化一个具体的 trainer 实例后，我们需要从 Trainer 的 + callback 类属性中将属于其的 callback 函数拿到，然后加入到 ``callback_manager`` 中； + + 这里的主要需要注意的地方在于为了支持没有带 ``marker`` 的 callback 函数赋给下方代码距离其最近的 trainer，在每次收集到 self._custom_callbacks[None] 后将其置为 []； """ _own_callbacks: List = copy.deepcopy(self._custom_callbacks["all"]) _own_callbacks.extend(self._custom_callbacks[None]) @@ -833,12 +846,24 @@ class Trainer(TrainerEventTrigger): self.add_callback_fn(*each_callback) def _check_callback_called_legality(self, check_mode: bool = True): - """ - 1. 函数的调用时机： + r""" + 这个函数主要的作用在于： + + 如果用户定制了训练流程中的一部分，例如 ``batch_step_fn`` 或者 ``TrainBatchLoop``；并且这些部分流程中可能会包含一些 callback + 函数的调用；例如 ``train_batch_loop.batch_step_fn`` 中包含 ``on_before_backward`` 等； + + 用户是十分可能忘记在其自己定制的部分流程中实现对这些 callback 函数的调用的；因此需要我们进行检测和提醒； + + 这种检测也十分简单，即如果我们检测到 callback_manager 的某一 callback 函数在训练一段时间（通常是涉及到允许定制的部分流程的结尾）后， + 其被调用的次数是 0，那么我们就会打印 ``warning`` 信息； + + 1. 这个函数的调用时机（这个函数会在以下情况被调用）： + 当检测 'batch_step_fn' 时，这个函数应当在 'train_batch_loop.run' 的 while 循环的最后进行调用； 当检测 'TrainBatchLoop' 时，这个函数应当在每一个 epoch 的最后进行调用； - 2. 函数作用 + 2. 这个函数作用的更细致的解释： + 这一函数的作用在于检查用户定制的 batch_step_fn / TrainBatchLoop 是否能够正确地调用 callback 函数，更准确地说，当用户实际 定制了 ("on_before_backward", "on_after_backward", "on_before_optimizers_step", "on_after_optimizers_step", "on_before_zero_grad", "on_after_zero_grad") / @@ -902,6 +927,9 @@ class Trainer(TrainerEventTrigger): """ Trainer 需要的一些 property """ @property def driver(self): + """ + :return: 返回 ``trainer`` 中的 ``driver`` 实例； + """ return self._driver @driver.setter @@ -910,6 +938,9 @@ class Trainer(TrainerEventTrigger): @property def train_batch_loop(self): + """ + :return: 返回 ``trainer`` 中的 ``train_batch_loop`` 实例； + """ return self._train_batch_loop @train_batch_loop.setter @@ -925,12 +956,24 @@ class Trainer(TrainerEventTrigger): def save_model(self, folder: Union[str, os.PathLike, BinaryIO, io.BytesIO], only_state_dict: bool = False, model_save_fn: Optional[Callable] = None, **kwargs): r""" - 用于帮助用户保存模型的辅助函数，具体实际的保存模型的操作由具体的 driver 实现； + 用于帮助您保存模型的辅助函数； + + :param folder: 保存模型的文件夹。如果没有传入 model_save_fn 参数，则我们会在这个文件夹下保存 fastnlp_model.pkl.tar 文件； + :param only_state_dict: 仅在 model_save_fn 为空时，有效。是否只保存模型的 ``state_dict``； + :param model_save_fn: 您自己定制的用来替换该保存函数本身保存逻辑的函数，当您传入了该参数后，我们会实际调用该函数，而不会去调用 ``driver`` 的 ``save_model`` 函数； + :param kwargs: 理论上您不需要使用到该参数； + + .. note:: - :param folder: 保存模型的文件夹。如果没有传入 model_save_fn 参数，则在这个文件夹下创建 fastnlp_model.pkl.tar 文件。 - :param only_state_dict: 仅在 model_save_fn 为空时，有效。是否只保存模型的 `state_dict`； - :param model_save_fn: 用户自己定制的用来替换该保存函数本身保存逻辑的函数； - :param kwargs: + 注意如果您需要在训练的过程中保存模型，如果没有特别复杂的逻辑，强烈您使用我们专门为保存模型以及断点重训功能定制的 ``callback``：**``CheckpointCallback``**； + ``CheckpointCallback`` 的使用具体见 :class:`fastNLP.core.callbacks.checkpoint_callback.CheckpointCallback`； + + 这意味着在大多数时刻您并不需要自己主动地调用该函数来保存模型；当然您可以在自己定制的 callback 类中通过直接调用 ``trainer.save_model`` 来保存模型； + + 具体实际的保存模型的操作由具体的 driver 实现，这意味着对于不同的 ``Driver`` 来说，保存模型的操作可能是不尽相同的， + 您如果想要了解更多的保存模型的细节，请直接查看各个 ``Driver`` 的 ``save_model`` 函数； + + ``save_model`` 函数和 ``load_model`` 函数是配套使用的； """ self.on_save_model() @@ -955,14 +998,22 @@ class Trainer(TrainerEventTrigger): def load_model(self, folder: Union[str, Path, BinaryIO, io.BytesIO], only_state_dict: bool = True, model_load_fn: Optional[Callable] = None, **kwargs): """ - 加载模型 - - :param folder: 读取 model 的文件夹，默认会尝试读取该文件夹下的 fastnlp_model.pkl.tar 文件。在 model_load_fn 不为空时， - 直接将该 folder 传递到 model_load_fn 中。 - :param only_state_dict: 要读取的文件中是否仅包含模型权重。在 model_load_fn 不为 None 时，该参数无意义。 - :param model_load_fn: callable 的函数，接受一个 folder 作为参数，不返回任何内容。 - :param kwargs: - :return: + 用于帮助您加载模型的辅助函数； + + :param folder: 存放着您需要加载的 model 的文件夹，默认会尝试读取该文件夹下的 fastnlp_model.pkl.tar 文件。在 model_load_fn 不为空时， + 直接将该 folder 传递到 model_load_fn 中； + :param only_state_dict: 要读取的文件中是否仅包含模型权重。在 ``model_load_fn 不为 None`` 时，该参数无意义； + :param model_load_fn: ``callable`` 的函数，接受一个 folder 作为参数，需要注意该函数不需要返回任何内容； + :param kwargs: 理论上您不需要使用到该参数； + + .. note:: + + 注意您需要在初始化 ``Trainer`` 后再通过 ``trainer`` 实例来调用该函数；这意味着您需要保证在保存和加载时使用的 ``driver`` 是属于同一个 + 训练框架的，例如都是 ``pytorch`` 或者 ``paddle``； + + 注意在大多数情况下您不需要使用该函数，如果您需要断点重训功能，您可以直接使用 ``trainer.load`` 函数； + + 该函数在通常情况下和 ``save_model`` 函数配套使用；其参数均与 ``save_model`` 函数成对应关系； """ self.on_load_model() self.driver.barrier() @@ -989,24 +1040,62 @@ class Trainer(TrainerEventTrigger): def save(self, folder: Union[str, Path], only_state_dict: bool = True, model_save_fn: Optional[Callable] = None, **kwargs): r""" - 用于断点重训 Trainer 的保存函数。 + 用于帮助您实现断点重训功能的保存函数； :param folder: 保存在哪个文件夹下，会在该文件下声称两个文件：fastnlp_checkpoint.pkl.tar 与 fastnlp_model.pkl.tar 。 - 如果 model_save_fn 不为空，则没有 fastnlp_model.pkl.tar 文件。 - :param only_state_dict: 当 model_save_fn 为空时有效，表明是否仅保存模型的权重。 - :param model_save_fn: 如果模型保存比较特殊，可以传入该函数自定义保存过程，输入应该接受一个文件夹（实际上就是接受上面的 folder - 参数），不必返回任何东西。 - :param kwargs: - :return: + 如果 model_save_fn 不为空，则没有 fastnlp_model.pkl.tar 文件； + :param only_state_dict: 当 model_save_fn 为空时有效，表明是否仅保存模型的权重； + :param model_save_fn: 如果模型保存比较特殊，可以传入该函数自定义模型的保存过程，输入应该接受一个文件夹（实际上就是接受上面的 folder + 参数），不需要返回值；这意味着您可以通过该函数来自己负责模型的保存过程，而我们则会将 ``trainer`` 的状态保存好； + :param kwargs: 理论上您不需要使用到该参数； + + .. note:: + + 注意如果您需要在训练的过程中使用断点重训功能，您可以直接使用 **``CheckpointCallback``**； + ``CheckpointCallback`` 的使用具体见 :class:`fastNLP.core.callbacks.checkpoint_callback.CheckpointCallback`； + + 这意味着在大多数时刻您并不需要自己主动地调用该函数来保存 ``Trainer`` 的状态；当然您可以在自己定制的 callback 类中通过直接调用 ``trainer.save`` 来保存 ``Trainer`` 的状态； + + 具体实际的保存状态的操作由具体的 driver 实现，这意味着对于不同的 ``Driver`` 来说，保存的操作可能是不尽相同的， + 您如果想要了解保存 ``Trainer`` 状态的更多细节，请直接查看各个 ``Driver`` 的 ``save`` 函数； + + ``save`` 函数和 ``load`` 函数是配套使用的； + + .. note:: + + 为了支持断点重训功能，我们会在调用该函数时保存以下内容： + + 1. 各个 ``callback`` 的状态，这主要涉及到一些带有运行状态的 ``callback``； + 2. 控制训练流程的变量 ``trainer_state``，具体详见 :class:`fastNLP.core.controllers.utils.states.TrainerState`； + 3. 一个特殊的变量 ``num_consumed_batches``，表示在这次训练过程中总共训练了多少个 batch 的数据；您不需要关心这个变量； + 4. sampler 的状态，为了支持断点重训功能，我们会在 trainer 初始化的时候，将您的 ``trainer_dataloader`` 的 ``sampler`` 替换为 + 我们专门用于断点重训功能的 ``ReproducibleSampler``，详见 :class:`fastNLP.core.samplers.reproducible_sampler.ReproducibleSampler`； + 5. model 的状态，即模型参数； + 6. optimizers 的状态，即优化器的状态； + 7. fp16 的状态； + + .. warning:: + + 一个值得注意的问题是 ``Driver`` 在新版 ``fastNLP`` 中的特殊作用，在断点重训时则体现为您应当尽量保证在前后两次训练中使用的 ``Driver`` + 是一致的，例如您不能在第一次训练时使用 ``pytorch``，而在第二次训练时使用 ``paddle``；或者尽量不要在第一次训练时使用分布式训练，但是 + 在第二次训练时使用非分布式训练（尽管这一行为的部分情况是支持的，请见下方的说明）； + + 但是如果您一定需要在前后使用不同分布式情况的 ``Driver``，那么在简单的默认情况下，我们也还是支持您使用断点重训的，这意味您可以在第一次训练时 + 使用单卡，但是在第二次训练时使用多卡进行训练；或者反过来； + + 以 ``pytorch`` 为例，这里的简单的默认情况指的是您的 ``train_dataloader`` 所使用的 ``sampler`` 是 ``RandomSampler`` 或者 ``SequentialSampler``； + 如果您的 ``sampler`` 是其它类型的 ``sampler``，那么我们仅支持前后两次训练 ``driver`` 严格不变时的断点重训； """ + self.driver.barrier() # 1. callback states 和 每一个callback的具体 callback 函数的 filter 的状态； # 2. trainer_state； - states = {"callback_states": self.on_save_checkpoint(), - "trainer_state": self.trainer_state.state_dict(), - 'num_consumed_batches': self.batch_idx_in_epoch - getattr(self, 'start_batch_idx_in_epoch', 0) - } + states = { + "callback_states": self.on_save_checkpoint(), + "trainer_state": self.trainer_state.state_dict(), + 'num_consumed_batches': self.batch_idx_in_epoch - getattr(self, 'start_batch_idx_in_epoch', 0) + } if isinstance(folder, str): folder = Path(folder) @@ -1025,18 +1114,34 @@ class Trainer(TrainerEventTrigger): def load(self, folder: str, resume_training: bool = True, only_state_dict: bool = True, model_load_fn: Optional[Callable] = None, **kwargs): r""" - 用于断点重训的加载函数； - 注意在 fastNLP 中断点重训的保存和加载逻辑是分开的，因此可能存在一种情况：用户只希望加载一个断点重训的状态，而在之后不再进行断点重训的 - 保存；在这种情况下，dataloader 的 sampler 就不一定会被替换成我们的 ReproducibleSampler； + 用于帮助您实现断点重训功能的加载函数； + + :param folder: 保存断点重训时 ``trainer`` 的状态文件的文件夹； + :param resume_training: 是否精确到从上次训练时最终截断的那一个 batch 开始训练；如果 ``resume_training=True``，那么我们 + 只会加载 ``model`` 和 ``optimizers`` 的状态；而其余对象的值则根据用户的 ``Trainer`` 的初始化直接重置； + :param only_state_dict: 保存的 ``model`` 是否只保存了权重； + :param model_load_fn: 使用的模型加载函数，参数应为一个文件夹，注意该函数不需要返回任何内容；您可以传入该参数来定制自己的加载模型的操作， + 当该参数不为 None 时，我们默认加载模型由该函数完成，``trainer.load`` 函数则会把 ``trainer`` 的其余状态加载好； + + .. note:: - 注意我们目前不支持单卡到多卡的断点重训； + 在 fastNLP 中，断点重训的保存和加载的逻辑是完全分离的，这意味着您在第二次训练时可以将 ``CheckpointCallback`` 从 ``trainer`` 中 + 去除，而直接使用 ``trainer.load`` 函数加载 ``trainer`` 的状态来进行断点重训； + + 该函数在通常情况下和 ``save`` 函数配套使用；其参数与 ``save`` 函数成对应关系； + + 对于在前后两次训练 ``Driver`` 不同的情况时使用断点重训，请参考 :meth:`fastNLP.core.controllers.trainer.Trainer.load` 函数的 ``warning``； + + Example:: + + trainer = Trainer(...) + + trainer.load(folder='/path-to-your-saved_checkpoint_folder/', ...) + + trainer.run() - :param folder: 保存断点重训 states 的文件地址； - :param resume_training: 是否从上次的 batch 开始训练，或者只从最近的 epoch 开始训练；注意如果 resume_training=True，那么我们 - 只会加载 model 和 optimizers 的状态；而其余的对象的值则根据用户的 Trainer 的初始化直接重置； - :param only_state_dict: 保存的 model 是否只包含了权重。 - :param model_load_fn: 使用的模型加载函数，参数应为一个 文件夹，不返回任何内容。 """ + self.driver.barrier() if isinstance(folder, str): folder = Path(folder) @@ -1083,12 +1188,30 @@ class Trainer(TrainerEventTrigger): """ 这四个函数是用来方便用户定制自己的 batch_step_fn（用于替换 train_batch_loop 当中的 batch_step_fn 函数） 的 """ def train_step(self, batch): + r""" + 实现模型训练过程中的对一个 batch 的数据的前向传播过程； + + .. note:: + + 该函数的提供是为了您能够更方便地定制自己的 ``train_batch_step_fn`` 来替换原本的 ``train_batch_loop.batch_step_fn``；更具体的细节 + 请见 :meth:`fastNLP.core.controllers.loops.train_batch_loop.TrainBatchLoop.batch_step_fn`； + + ``trainer.backward / zero_grad / step`` 函数的作用类似； + + :param batch: 一个 batch 的数据； + :return: 返回模型的前向传播函数所返回的结果； + """ with self.driver.auto_cast(): outputs = self.driver.model_call(batch, self._train_step, self._train_step_signature_fn) outputs = match_and_substitute_params(self.output_mapping, outputs) return outputs def backward(self, outputs): + r""" + 实现模型训练过程中神经网络的反向传播过程； + + :param outputs: 模型的输出，应当为一个字典或者 dataclass，里面包含以 ``loss`` 为关键字的值； + """ self.on_before_backward(outputs) loss = self.extract_loss_from_outputs(outputs) loss = loss / self.accumulation_steps @@ -1096,27 +1219,40 @@ class Trainer(TrainerEventTrigger): self.on_after_backward() def zero_grad(self): + r""" + 实现模型训练过程中对优化器中的梯度的置零操作； + """ if (self.global_forward_batches + 1) % self.accumulation_steps == 0: self.on_before_zero_grad(self.optimizers) self.driver.zero_grad(self.set_grad_to_none) self.on_after_zero_grad(self.optimizers) def step(self): + r""" + 实现模型训练过程中的优化器的参数更新操作； + """ + if (self.global_forward_batches + 1) % self.accumulation_steps == 0: self.on_before_optimizers_step(self.optimizers) self.driver.step() self.on_after_optimizers_step(self.optimizers) def move_data_to_device(self, batch): + r""" + 将数据迁移到当前进程所使用的设备上； + + :param batch: 一个 batch 的数据； + :return: 位置已经被迁移后的数据； + """ return self.driver.move_data_to_device(batch) @staticmethod def extract_loss_from_outputs(outputs): r""" - 用来从用户模型的输出对象中抽取 `loss` 对象； - 目前支持 `outputs` 对象为 'Dict' 或者 'dataclass'； + 用来从用户模型的输出对象中抽取 ``loss`` 对象； + 目前支持 `outputs` 对象为 ``dict`` 或者 ``dataclass``； - :return: 返回被抽取出来的 `loss` 对象，如果当前运行的是 'pytorch' 的 `Driver`，那么返回的就是一个 tensor； + :return: 返回被抽取出来的 ``loss`` 对象，例如如果是 ``pytorch``，那么返回的就是一个 tensor； """ if isinstance(outputs, Dict): try: @@ -1139,10 +1275,10 @@ class Trainer(TrainerEventTrigger): @contextmanager def get_no_sync_context(self): r""" - 用于在梯度累积并且使用 DDP 时，由于在前 `accumulation_steps` - 1 的时间内不需要进行梯度的同步，因此通过使用该 context 上下文 - 环境来避免梯度的同步； + 用于在使用梯度累积并且进行分布式训练时，由于在前 ``accumulation_steps - 1`` 的时间内不需要进行梯度的同步，因此通过使用该 context 上下文 + 环境来避免梯度的同步； - :return: 一个 no_sync 的 context； + :return: 一个支持 ``no_sync`` 的 ``context``； """ if (self.global_forward_batches + 1) % self.accumulation_steps != 0: @@ -1157,6 +1293,9 @@ class Trainer(TrainerEventTrigger): @property def n_epochs(self) -> int: + r""" + :return: 返回当前训练的总体的 epoch 的数量； + """ return self.trainer_state.n_epochs @n_epochs.setter @@ -1165,6 +1304,9 @@ class Trainer(TrainerEventTrigger): @property def cur_epoch_idx(self) -> int: + r""" + :return: 返回当前正在第几个 epoch； + """ return self.trainer_state.cur_epoch_idx @cur_epoch_idx.setter @@ -1173,6 +1315,9 @@ class Trainer(TrainerEventTrigger): @property def global_forward_batches(self) -> int: + """ + :return: 返回从训练开始到当前总共训练了多少 batch 的数据； + """ return self.trainer_state.global_forward_batches @global_forward_batches.setter @@ -1181,6 +1326,9 @@ class Trainer(TrainerEventTrigger): @property def batch_idx_in_epoch(self) -> int: + r""" + :return: 返回在从当前的这个 epoch 开始，到现在共训练了多少 batch 的数据； + """ return self.trainer_state.batch_idx_in_epoch @batch_idx_in_epoch.setter @@ -1189,6 +1337,9 @@ class Trainer(TrainerEventTrigger): @property def num_batches_per_epoch(self) -> int: + r""" + :return: 返回每一个 epoch 实际会训练多少个 batch 的数据； + """ return self.trainer_state.num_batches_per_epoch @num_batches_per_epoch.setter @@ -1197,6 +1348,9 @@ class Trainer(TrainerEventTrigger): @property def total_batches(self) -> int: + r""" + :return: 返回整体的训练中实际会训练多少个 batch 的数据； + """ return self.trainer_state.total_batches @total_batches.setter @@ -1207,16 +1361,27 @@ class Trainer(TrainerEventTrigger): @property def model_device(self): + r""" + :return: 返回当前模型所在的设备；注意该值在当且仅当在少数情况下为 ``None``，例如当使用 ``pytorch`` 时，仅当用户自己初始化 ``init_progress_group`` 时 + ``model_device`` 才为 None； + """ return self.driver.model_device @property def data_device(self): + r""" + :return: 返回数据会被迁移到的目的设备； + """ return self.driver.data_device """ dataloader property """ @property def train_dataloader(self): + """ + :return: 返回用户传入的 ``train_dataloader``，注意该 ``dataloader`` 与用户传入给 ``Trainer`` 的 ``dataloader`` 对象是同一个对象，而我们在 + 实际训练过程中使用的 ``dataloader`` 的状态可能有所更改； + """ return self._train_dataloader @train_dataloader.setter @@ -1225,6 +1390,9 @@ class Trainer(TrainerEventTrigger): @property def evaluate_dataloaders(self): + """ + :return: 返回用户传入的 ``evaluate_dataloaders``； + """ return self._evaluate_dataloaders @evaluate_dataloaders.setter @@ -1234,6 +1402,10 @@ class Trainer(TrainerEventTrigger): def _get_input_output_mapping(input_mapping, output_mapping, train_input_mapping, train_output_mapping, evaluate_input_mapping, evaluate_output_mapping): + """ + 确定在训练过程中到底要使用哪个 input_mapping 和 output_mapping，之所以要设置该函数是因为在有些时候 evaluate 所需要的 input_mapping 和 + output_mapping 是与 train 的时候是不一样的，因此需要额外的定制； + """ if train_input_mapping is not None and input_mapping is not None: raise ValueError("Parameter `input_mapping` and `train_input_mapping` cannot be set simultaneously.") diff --git a/fastNLP/core/drivers/jittor_driver/jittor_driver.py b/fastNLP/core/drivers/jittor_driver/jittor_driver.py index 7efff348..25fc4af6 100644 --- a/fastNLP/core/drivers/jittor_driver/jittor_driver.py +++ b/fastNLP/core/drivers/jittor_driver/jittor_driver.py @@ -21,6 +21,9 @@ if _NEED_IMPORT_JITTOR: 'sum': jt.sum } +__all__ = [ + "JittorDriver", +] class JittorDriver(Driver): r""" @@ -90,9 +93,6 @@ class JittorDriver(Driver): "'test_step'.") def save_model(self, filepath: str, only_state_dict: bool = False, model_save_fn: Optional[Callable]=None): - """ - 保存模型 - """ if model_save_fn is not None: outputs = model_save_fn(filepath) if outputs is not None: @@ -105,12 +105,6 @@ class JittorDriver(Driver): jt.save(states, filepath) def load_model(self, filepath: str): - """ - 加载模型的加载函数； - - :param file_path: 保存文件的文件位置（需要包括文件名）； - :return: 加载后的state_dict - """ if not os.path.exists(filepath): raise FileNotFoundError("Checkpoint at {} not found.".format(filepath)) return jt.load(filepath) @@ -156,7 +150,7 @@ class JittorDriver(Driver): def move_data_to_device(self, batch: 'jt.Var'): """ - jittor暂时没有提供数据迁移的函数，因此这个函数只是简单地返回batch + **jittor** 暂时没有提供数据迁移的函数，因此这个函数只是简单地返回 **batch** """ return batch diff --git a/fastNLP/core/drivers/jittor_driver/mpi.py b/fastNLP/core/drivers/jittor_driver/mpi.py index bfa49e68..4ade3fd1 100644 --- a/fastNLP/core/drivers/jittor_driver/mpi.py +++ b/fastNLP/core/drivers/jittor_driver/mpi.py @@ -20,6 +20,10 @@ class JittorMPIDriver(JittorDriver): 这是一个正在开发中的功能，敬请期待。 + .. todo: + + 实现断点重训中替换 dataloader 的 set_dist_repro_dataloader 函数 + """ def __init__( self, diff --git a/fastNLP/core/drivers/jittor_driver/utils.py b/fastNLP/core/drivers/jittor_driver/utils.py index 43be9ac3..c6d44cfc 100644 --- a/fastNLP/core/drivers/jittor_driver/utils.py +++ b/fastNLP/core/drivers/jittor_driver/utils.py @@ -9,7 +9,7 @@ __all__ = [] class DummyGradScaler: """ - 用于仿造的GradScaler对象，防止重复写大量的if判断 + 用于仿造的 **GradScaler** 对象，防止重复写大量的if判断 """ def __init__(self, *args, **kwargs): pass diff --git a/fastNLP/core/drivers/paddle_driver/dist_utils.py b/fastNLP/core/drivers/paddle_driver/dist_utils.py index ffa142d3..4dea268d 100644 --- a/fastNLP/core/drivers/paddle_driver/dist_utils.py +++ b/fastNLP/core/drivers/paddle_driver/dist_utils.py @@ -21,6 +21,8 @@ if _NEED_IMPORT_PADDLE: _parse_load_result, ) +__all__ = [] + def _validate_output_list_for_rank(my_rank, dst, gather_list): if dst == my_rank: if not gather_list: diff --git a/fastNLP/core/drivers/paddle_driver/fleet.py b/fastNLP/core/drivers/paddle_driver/fleet.py index 03dc3375..36fb74fd 100644 --- a/fastNLP/core/drivers/paddle_driver/fleet.py +++ b/fastNLP/core/drivers/paddle_driver/fleet.py @@ -1,3 +1,69 @@ +r""" +用于实现 **PaddlePaddle** 框架下使用 ``fleet`` 分布式训练 API 进行集群式（*collective*）多卡训练的 Driver。 + +.. note:: + + 在 **PaddlePaddle** 框架中，使用分布式训练的方式可以参见 **PaddlePaddle** 的 + `官方文档 `_ 。 + 简言之，分布式训练的过程可以概括为：导入 ``fleet`` 包 -> 使用 :func:`fleet.init` 初始化分布式环境 -> 初始化模型，转换为并行模型开始训练。 + +**fastNLP** 支持三种启动分布式训练的方式（假设执行训练的文件名为 ``train.py``）： + + A. 用户自己不进行分布式的任何操作，直接使用我们的 :class:`~fastNLP.core.Trainer` 进行训练，此时将参数 ``device`` + 设置为一个列表，然后使用 ``python train.py`` 的方式开始训练； + B. 用户自己不进行分布式的任何操作，但是使用 ``python -m paddle.distributed.launch train.py`` 开始训练； + C. 用户自己在外面初始化分布式环境，并且通过 ``python -m paddle.distributed.launch train.py`` 开始训练； + +.. note:: + + 在后两种启动方式中，您需要通过参数 ``--gpus`` 来指定训练使用的设备，在 ``trainer`` 中设置的参数是无效的。 + +不过在使用该 Driver 之前，我们需要向您说明 **fastNLP** 实现 ``PaddleFleetDriver`` 的思路，以便于您理解代码编写过程中可能出现的问题。 + +在 **fastNLP** 中，为了尽可能减少单卡向分布式训练转换过程中的代码变动，我们需要在 ``PaddleFleetDriver`` 中进行 **分布式环境初始化** +和 **将模型转换为并行模式** 等操作，同时实现多卡训练的方法是从主进程（``rank=0``）中创建其它的所有子进程（``rank=1,2,...``）。 +在这个过程中，我们发现由于 **PaddlePaddle** 框架的特性，会出现下面的问题： + + 1. **fastNLP** 中，初始化模型一定会在初始化 ``Driver`` 之前，因此调用 :func:`fleet.init` 的时机会在初始化模型之后； + 此时子进程中模型将无法正常地初始化，提示无法找到设备 ``gpu:0``； + 2. 在训练的过程中，会出现训练一个 ``batch`` 后程序卡住或程序会占用所有可见显卡的情况； + +考虑到这些问题，我们为 **PaddlePaddle** 的分布式训练制定了这样的约束：在导入 **fastNLP** 之前，必须设置环境变量 ``FASTNLP_BACKEND`` +为 ``paddle``。执行方法有两种:: + + >>> import os + >>> os.environ["FASTNLP_BACKEND"] = "paddle" # 设置环境变量 + >>> import fastNLP # 设置之后才可以导入 fastNLP + +或是在执行脚本（假设文件名为 ``train.py`` ）时设置:: + + FASTNLP_BACKEND=paddle python train.py + FASTNLP_BACKEND=paddle python -m paddle.distributed.lauch train.py + +设置 ``FASTNLP_BACKEND=paddle`` 后，**fastNLP** 会在 ``import paddle`` 之前通过 ``CUDA_VISIBLE_DEVICES`` 将设备限制在所有可见设备的第 +**0** 张卡上，以此绕开通信和同步上的种种限制。我们会将用户希望可见的设备（如用户自己设置了 ``CUDA_VISIBLE_DEVICES`` 的情况）保存在另一个环境变量 +``USER_CUDA_VISIBLE_DEVICES`` 中来确保 **fastNLP** 能够知道用户的设置。假设用户希望在 ``[0,2,3]`` 三张显卡上进行分布式训练，那么在三个训练进程中， +``CUDA_VISIBLE_DEVICES`` 就分别为 0、2 和 3 。 + +.. note:: + + 我们会事先将设备限制在所有可见设备的第 **0** 张卡上，因此多卡训练的参数 ``device`` 一定要以 **0** 开始，否则会无法正常地启动。 + 如果您希望调整使用的第一张显卡，请使用 ``CUDA_VISIBLE_DEVICES`` 进行限制。 + +.. note:: + + 根据 **PaddlePaddle** 的说明，设置 ``CUDA_VISIBLE_DEVICES`` 之后启动分布式训练时，情况A与情况BC设置设备的方式会有所不同。 + 情况A应设置为实际设备相对可见设备的索引，而情况BC应设置为实际的设备号： + + 1. 情况A中， ``CUDA_VISIBLE_DEVICES=3,4,5,6`` 且参数 ``device=[0,2,3]`` 代表使用 **3号、5号和6号** 显卡； + 2. 情况BC中，``CUDA_VISIBLE_DEVICES=3,4,5,6`` 且参数 ``--gpu=3,5,6`` 代表使用 **3号、5号和6号** 显卡； + +.. note:: + + 多机的启动强制要求用户在每一台机器上使用 ``python -m paddle.distributed.launch`` 启动；因此我们不会在 ``PaddleFleetDriver`` + 中保存任何当前有多少台机器的信息； + +""" import os from typing import List, Union, Optional, Dict, Tuple, Callable @@ -53,6 +119,33 @@ __all__ = [ ] class PaddleFleetDriver(PaddleDriver): + """ + :param model: 训练使用的模型； + + * 如果不想自己初始化分布式环境，类型应为 :class:`paddle.nn.Layer`； + * 如果已经在外面初始化了分布式环境，类型应为 :class:`paddle.DataParallel`； + + :param parallel_device: 多卡训练时使用的设备，必须是一个列表。 + 当使用 ``python -m paddle.distributed.launch`` 启动时，该参数无效； + :param is_pull_by_paddle_run: 标记当前进程是否为通过 ``python -m paddle.distributed.launch`` 启动的。 + 这个参数仅在 :class:`~fastNLP.core.Trainer` 中初始化 driver 时使用 + :param fp16: 是否开启混合精度训练； + :kwargs: + * *paddle_kwargs* -- 用于在指定 ``driver`` 为 'paddle' 时设定具体 driver 实例的一些参数： + + * fleet_kwargs -- 用于在使用 ``PaddleFleetDriver`` 时指定 ``DataParallel`` 和 ``fleet`` 初始化时的参数，包括： + + * is_collective -- 是否使用 paddle 集群式的分布式训练方法，目前仅支持为 ``True`` 的情况； + * role_maker -- 初始化 ``fleet`` 分布式训练 API 时使用的 ``RoleMaker`` + * 其它用于初始化 ``DataParallel`` 的参数； + + * 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, @@ -61,143 +154,20 @@ class PaddleFleetDriver(PaddleDriver): fp16: bool = False, **kwargs ): - r""" - 通过使用 PaddlePaddle 的 Fleet 框架启动多卡进程的 Driver。 - 需要注意的一点是，由于 PaddlePaddle 框架的特性，如果直接使用在 rank0 拉起其它进程的方法的话，如果不加以任何限制，PaddlePaddle会出现 - 第一次前向传播后卡住或占用所有显卡的现象；为了解决这一问题，我们在引入 FastNLP 时，会使用 `CUDA_VISIBLE_DEVICES` 将设备限制在卡0上， - 而用户如果使用了这一环境变量，我们会将其储存在 `USER_CUDA_VISIBLE_DEVICES` 中，并且通过一定的手段实现了转换（详细的设置请参见： - `fastNLP/envs/set_backend.py`）。在拉起其它进程的时候，我们会如法炮制，将环境限制在对应的设备上。 - - `PaddleFleetDriver` 目前支持的三种启动方式： - 1. 用户自己不进行分布式的任何操作，直接使用我们的 Trainer，这时是由我们自己使用 `FleetLauncher` 拉起多个进程， - 然后 `PaddleFleetDriver` 自己通过调用 `fleet.init` 来初始化 ddp 的通信组；（情况 A） - 2. 用户同样不在 Trainer 之外初始化分布式训练，但是用户自己使用 python -m paddle.distributed.launch 拉起来创建多个进程，这时我们仍旧 - 会通过调用 `fleet.init` 来初始化 ddp 的通信组；（情况 B） - 3. 用户自己在外面初始化分布式，并且通过 python -m paddle.distributed.launch 拉起，这时无论是多个进程的拉起和通信组的建立 - 都由用户自己操作，我们只会在 driver.setup 的时候对 `PaddleFleetDriver` 设置一些必要的属性值；（情况 C） - - 注意多机的启动强制要求用户在每一台机器上使用 python -m paddle.distributed.launch 启动；因此我们不会在 `PaddleFleetDriver` 中保存 - 任何当前有多少台机器的信息； - - Part 1：三种启动方式的具体分析： - （1）对于用户运行的脚本中，如果 `driver.setup` 只会被调用一次（意味着用户的启动脚本中只初始化了一个 trainer/evaluator）时， - `PaddleFleetDriver` 在初始化以及 `setup` 函数中会做的事情分别如下所示： - -> 情况 A：这种情况下用户传入的 model 在一定是普通的 model（没有经 `DataParallel` 包裹的model）， - 因为 `Parallel` 的使用一定要求 fleet.init 已经被调用用来建立当前的 ddp 通信组；但是这意味着如果 - 用户需要使用 2 张以上的显卡，那么其必然需要使用 paddle.distributed.launch 来启动，意味着就不是情况 A 了； - 这时我们首先会调用 `FleetLauncher.launch` 函数来拉起多个进程，其中进程的数量等于用户传入给 trainer 的使用的 gpu - 的数量（例如 `Trainer` 中的参数是 device=[0, 1, 6, 7]，那么我们就会使用第 0、1、6、7 张 gpu 来拉起 4 个进程）； - 接着我们会调用 `fleet.init` 来初始化各个进程之间的通信组； - 这里需要注意拉起的新的进程会从前到后完整地运行一遍用户的启动脚本（例如 main.py），因此也都会运行这两个函数，但是需要注意只有进程 0 - 才会去真正地运行 `FleetLauncher.launch`；进程 0 运行到 `fleet.init`，paddle 会阻塞进程 0 继续 - 向前运行，直到其它进程也运行到这里； - 最后我们会设置这个进程对应的 device，然后将模型迁移到对应的机器上，再使用 `DataParallel` 将模型包裹； - 至此，paddle 分布式的环境配置过程全部完成； - - -> 情况 B：注意这种情况我们直接限定了用户是通过 paddle.distributed.launch 拉起，并且没有自己建立分布式的通信组。这时在 - `PaddleFleetDriver` 的初始化和 setup 函数的调用过程中，与情况 A 首要的不同就在于用户在 trainer 中输入的参数 device 不再有效， - 这时每个进程所使用的 gpu 是我们直接通过 `CUDA_VISIBLE_DEVICE` 来配置的；因此，如果用户想要实现使用特定 gpu - 设备的目的，可以通过自己设置环境变量实现（例如 os.environ["CUDA_VISIBLE_DEVICE"] 来实现，我们会通过一定的手段将其保存起来）； - 剩下的操作和情况 A 类似； - - -> 情况 C：注意这种情况我们限定了用户是通过 paddle.distributed.launch 拉起，并且 ddp 的通信组也是由自己建立。这时基本上所有的 - 与操作相关的操作都应当由用户自己完成，包括迁移模型到对应 gpu 上以及将模型用 `DataParallel` 包裹等。 - （2）如果 `driver.setup` 函数在脚本中会被调用两次及以上（意味着用户的启动脚本初始化了两个及以上的 trainer/evaluator）时： - 注意这种情况下我们是会保证前后两个 trainer/evaluator 使用的 `PaddleFleetDriver` 以及其初始化方式的一致性，换句话说，如果 trainer1 - 检测到的启动方式是 '情况 A'，那么我们会保证 trainer2 检测到的启动方式同样是 '情况A'（即使这需要一些额外的处理）；因此这里我们主要讨论 - 我们是通过怎样的操作来保证 trainer2/3/... 检测到的启动方式是和 trainer1 一致的；简单来说，我们是通过使用环境变量来标记每一种不同的 - 启动方式来实现这一点的： - 我们会使用 `FASTNLP_DISTRIBUTED_CHECK` 来标记 '情况 A'，使用 `fastnlp_torch_launch_not_ddp` 来标记 '情况 B'，意味着我们在 - 使用 '情况 A' 来启动 `PaddleFleetDriver` 时，我们会将 `FASTNLP_DISTRIBUTED_CHECK` 这一字符串注入到环境变量中，而 '情况 B' 时则 - 会将 `fastnlp_torch_launch_not_ddp` 这一字符串注入到环境变量中。因此在 trainer2 的 `PaddleFleetDriver` 的初始化和 setup 过程中， - 如果检测到这些特殊的环境变量，我们就会将启动方式变更为其对应的启动方式，即使其它的参数特征属于另外的启动方式。 - - Part 2：对应的代码细节： - 1. 如何判断当前的各进程之间的通信组已经被建立（fleet 已经被初始化）； - parallel_helper._is_parallel_ctx_initialized()； - 2. 如何判断不同的进程是否是由 `python -m paddle.distributed.launch` 拉起还是由我们的 `FleetLauncher.launch()` - 函数拉起； - 我们会在用户脚本 `import fastNLP` 的时候检测当前的环境变量中是否有 'PADDLE_RANK_IN_NODE'、'PADDLE_TRAINER_ID' - 以及没有 `FASTNLP_DISTRIBUTED_CHECK`， - 如果满足条件，则我们会向环境变量中注入特殊的值 'FASTNLP_BACKEND_LAUNCH' 来标记用户是否使用了 `python -m paddle.distributed.launch` - 来拉起多个进程； - 3. 整体的处理判断流程： - ___________________________________ - ｜进入 PaddleFleetDriver 的 __init__ 函数｜ - ——————————————————————————————————— - ↓ - ___________________________________________________ - ｜ 判断不同的进程是否是由 paddle.distributed.launch 拉起 ｜ - ｜（或者我们自己的 FleetLauncher 函数拉起） ｜ --------------> - ——————————————————————————————————————————————————— 　｜ - ↓ 是由 paddle.distributed.launch 拉起 ｜ 我们自己的 FleetLauncher 函数拉起多个进程 - 　_____________________________ 　　　　　　　　　　　｜　 - ←←←←← ｜ 检测用户是否自己初始化了 fleet ｜ 　　　　　　　　　　　　　｜ - ↓ —————————————————————————————　　　　　　　　　　　　　　　　　 ↓ - ↓ ↓ 是 ________ - ↓ ______ ｜ 情况 A ｜ - ↓ 否 |情况 C| ————————— - ↓ ——————— - ↓ - ↓ ______ - ↓ -----------> |情况 B| - 　　——————— - 4. 为了完成全部的建立分布式所需要的操作，三种情况都需要做的事情，以及每件事情的职责归属： - - 情况 A ｜ 情况 B ｜ 情况 C - ________________________________________________________________________________________________________ - 配置 fleet 所 ｜ FleetLauncher.launch ｜ paddle.distributed.launch｜ paddle.distributed.launch - 需要的环境变量 ｜ ｜ ｜ - ———————————————————————————————————————————————————————————————————————————————————————————————————————— - 开启多个进程 ｜ FleetLauncher.launch ｜ paddle.distributed.launch｜ paddle.distributed.launch - ———————————————————————————————————————————————————————————————————————————————————————————————————————— - 调用 fleet.init函数 ｜ PaddleFleetDriver.setup ｜ PaddleFleetDriver.setup ｜ 用户自己调用 - ———————————————————————————————————————————————————————————————————————————————————————————————————————— - 设置 PaddleFleetDriver ｜ ｜ ｜ - 的 world_size 和 ｜ PaddleFleetDriver.setup ｜ PaddleFleetDriver.setup ｜ PaddleFleetDriver.setup - global_rank 属性 ｜ ｜ ｜ - ———————————————————————————————————————————————————————————————————————————————————————————————————————— - - Part 3：其它的处理细节： - 1. 环境变量； - fastNLP 的 `PaddleFleetDriver` 运行时所需要的环境变量分为两种，一种是 paddle fleet 运行所需要的环境变量；另一种是 fastNLP 自己 - 的环境变量。前者的配置情况如上表所示；而后者中的大多数环境变量则是在用户 import fastNLP 时就设置好了； - 2. parallel_device, model_device 和 data_device 的关系； - parallel_device 为 `PaddleFleetDriver` 的参数，model_device 和 data_device 都为 driver 的属性； - 其中 data_device 仅当情况 C 时由用户自己指定；如果其不为 None，那么在模型 forward 的时候，我们就会将数据迁移到 data_device 上； - model_device 永远都为单独的一个 torch.device； - - 情况 A ｜ 情况 B ｜ 情况 C - ________________________________________________________________________________________________________ - parallel_device ｜ 由用户传入trainer的参数 ｜ ｜ - ｜ device 决定，必须是一个list， ｜ 为 CUDA_VISIBLE_DEVICES ｜ 为 CUDA_VISIBLE_DEVICES - ｜ 其中每一个对象都是 int ｜ ｜ - ———————————————————————————————————————————————————————————————————————————————————————————————————————— - model_device ｜ parallel_device[local_rank] ｜ parallel_device ｜ None - ———————————————————————————————————————————————————————————————————————————————————————————————————————— - data_device ｜ model_device ｜ model_device ｜ 由用户传入 trainer 的参数 - ｜ ｜ ｜ data_device 决定 - ———————————————————————————————————————————————————————————————————————————————————————————————————————— - - 3. _DDPWrappingModel 的作用； - 因为我们即需要调用模型的 `train_step`、`evaluate_step`、`test_step` 方法，又需要通过 `DataParallel` 的forward 函数来帮助 - 我们同步各个设备上的梯度，因此我们需要先将模型单独包裹一层，然后在 forward 的时候，其先经过 `DataParallel` 的 forward 方法， - 然后再经过 `_DDPWrappingModel` 的 forward 方法，我们会在该 forward 函数中进行判断，确定调用的是模型自己的 forward 函数，还是 - `train_step`、`evaluate_step`、`test_step` 方法。 - - 4. 当某一个进程出现 exception 后，`PaddleFleetDriver` 的处理； - - 不管是什么情况，`PaddleFleetDriver` 在 `setup` 函数的最后，都会将所有进程的 pid 主动记录下来，这样当一个进程出现 exception 后， - driver 的 on_exception 函数就会被 trainer 调用，其会调用 os.kill 指令将其它进程 kill 掉； - """ if USER_CUDA_VISIBLE_DEVICES not in os.environ: raise RuntimeError("To run paddle distributed training, please set `FASTNLP_BACKEND` to 'paddle' before using FastNLP.") super(PaddleFleetDriver, self).__init__(model, fp16=fp16, **kwargs) # 如果不是通过 launch 启动，要求用户必须传入 parallel_device - if not is_pull_by_paddle_run and parallel_device is None: - raise ValueError("Parameter `parallel_device` can not be None when using `PaddleFleetDriver`. This error is caused " - "when your value of parameter `device` is `None` in your `Trainer` instance.") + if not is_pull_by_paddle_run: + if parallel_device is None: + raise ValueError("Parameter `parallel_device` can not be None when using `PaddleFleetDriver`. This error is caused " + "when your value of parameter `device` is `None` in your `Trainer` instance.") + if not isinstance(parallel_device, List): + raise ValueError("Parameter `parallel_device`'s type must be List when using `PaddleFleetDriver`, " + f"not {type(parallel_device)}.") + if get_paddle_device_id(parallel_device[0]) != 0: + raise ValueError("The first device of `parallel_device` must be 'gpu:0' in fastNLP.") # 如果用户自己初始化了 paddle 的分布式训练那么一定是通过 launch 拉起的 # 这个参数会在 initialize_paddle_drvier 中设置。 @@ -254,10 +224,10 @@ class PaddleFleetDriver(PaddleDriver): def setup(self): """ - 根据不同的情况进行不同的设置。 - 1、如果是通过 paddle.distributed.launch 方法启动时，则根据已经设置好的环境获取 - 分布式的属性。 - 2、否则，调用 FleetLauncher 类启动子进程 + 初始化分布式训练的环境。 + + 1. 如果是通过 ``paddle.distributed.launch`` 方法启动的，则根据已经设置好的环境获取分布式的属性。 + 2. 否则启动子进程。 """ if self._has_setup: return @@ -267,7 +237,7 @@ class PaddleFleetDriver(PaddleDriver): if self.outside_fleet: # 已经初始化了多机环境 - self.set_from_fleet_environment() + self._set_from_fleet_environment() else: # 用户没有初始化多机环境 # TODO 绕一下 @@ -287,7 +257,7 @@ class PaddleFleetDriver(PaddleDriver): # parallel_device 是 list， if not parallel_helper._is_parallel_ctx_initialized(): # 拉起子进程并设置相应的属性 - self.init_fleet_and_set() + self._init_fleet_and_set() # 用户在这个 trainer 前面又初始化了一个 trainer，并且使用的是 PaddleFleetDriver； else: # 已经设置过一次，保证参数必须是一样的 @@ -321,7 +291,7 @@ class PaddleFleetDriver(PaddleDriver): self._pids = self._pids[node_rank*local_world_size: (node_rank+1)*local_world_size] self._pids = self.tensor_to_numeric(self._pids) - def init_fleet_and_set(self): + def _init_fleet_and_set(self): """ 使用 FleetLauncher 拉起子进程 """ @@ -340,7 +310,7 @@ class PaddleFleetDriver(PaddleDriver): assert self.world_size is not None assert self.world_size == len(self.parallel_device) - def set_from_fleet_environment(self): + def _set_from_fleet_environment(self): """ 当用户使用了 `python -m paddle.distributed.launch xxx.py` 启动时，我们需要 根据 paddle 设置的环境变量来获得各种属性 @@ -349,19 +319,11 @@ class PaddleFleetDriver(PaddleDriver): self.global_rank = paddledist.get_rank() def barrier(self): - r""" - 用于在多进程工作时同步各进程的工作进度，运行快的进程运行到这里会等待运行慢的进程，只有所有进程都运行到此函数时，所有的进程才会继续运行； - 仅在多分布式训练场景中有使用。 - - 注意，该函数的行为会受到 FASTNLP_NO_SYNC 的影响。仅当 FASTNLP_NO_SYNC 在 os.environ 中不存在，或小于 1 时才真的执行 barrier 。 - """ if int(os.environ.get(FASTNLP_NO_SYNC, 0)) < 1: # 当 FASTNLP_NO_SYNC 小于 1 时实际执行 paddledist.barrier() def configure_fleet(self): - """ - 将模型用 DataParallel 和自定义的类型包裹起来 - """ + # 将模型用 DataParallel 和自定义的类型包裹起来 if not self._has_fleetwrapped and not isinstance(self.model, DataParallel): self.model = DataParallel( _FleetWrappingModel(self.model), @@ -395,10 +357,17 @@ class PaddleFleetDriver(PaddleDriver): @property def model_device(self): + """ + :return: 模型所在的设备； + """ return self._model_device @property def data_device(self): + """ + :return: 数据所在的设备；由于 **PaddlePaddle** 可以通过环境变量获取当前进程的设备，因此该属性 + 和 ``model_device`` 表现相同； + """ return self.model_device def model_call(self, batch, fn: Callable, signature_fn: Optional[Callable]) -> Dict: @@ -522,23 +491,29 @@ class PaddleFleetDriver(PaddleDriver): else: raise ValueError("Parameter `dist_sampler` can only be one of three values: ('dist', 'unrepeatdist', None).") - def is_global_zero(self): + def is_global_zero(self) -> bool: return self.global_rank == 0 def get_model_no_sync_context(self): return self.model.no_sync - def unwrap_model(self): + def unwrap_model(self) -> "paddle.nn.Layer": + """ + 获得 driver 最原始的模型。该函数可以取出被 :class:`paddle.DataParallel` 包裹起来的模型。 + """ _layers = self.model._layers if isinstance(_layers, _FleetWrappingModel): return _layers.model else: return _layers - def get_local_rank(self) ->int: + def get_local_rank(self) -> int: return self.local_rank - def is_distributed(self): + def is_distributed(self) -> bool: + """ + 判断是否为分布式的 **Driver** ，在 ``PaddleFleetDriver`` 中，返回 ``True``。 + """ return True @staticmethod diff --git a/fastNLP/core/drivers/paddle_driver/paddle_driver.py b/fastNLP/core/drivers/paddle_driver/paddle_driver.py index 74c7b7a8..c10c98a1 100644 --- a/fastNLP/core/drivers/paddle_driver/paddle_driver.py +++ b/fastNLP/core/drivers/paddle_driver/paddle_driver.py @@ -1,14 +1,12 @@ import os import random -from typing import Union, Optional, Dict +from typing import Union, Optional, Dict, Any from pathlib import Path from functools import partial from dataclasses import dataclass import numpy as np -from fastNLP.envs.env import USER_CUDA_VISIBLE_DEVICES - from .utils import _build_fp16_env, optimizer_state_to_device, DummyGradScaler from fastNLP.envs.imports import _NEED_IMPORT_PADDLE from fastNLP.core.drivers.driver import Driver @@ -50,9 +48,26 @@ if _NEED_IMPORT_PADDLE: class PaddleDriver(Driver): r""" - Paddle框架的Driver，包括实现单卡训练的 `PaddleSingleDriver` 和分布式训练的 `PaddleFleetDriver`。 + 实现了 **PaddlePaddle** 框架训练功能的基本 Driver，实现了单卡和多卡情景下均需要实现的功能，以和 **fastNLP** 的 + :class:`~fastNLP.core.Trainer` 兼容；通过这个 Driver，可以在 **fastNLP** 中实现从 **Pytorch** 框架到 + **PaddlePaddle** 深度学习框架的切换。 + + 这个类被以下子类继承： + + 1. :class:`~fastNLP.core.drivers.PaddleSingleDriver`：实现了使用单卡和 ``cpu`` 训练的具体功能； + 2. :class:`~fastNLP.core.drivers.PaddleFleetDriver`：实现了使用 ``fleet`` 分布式训练 API 进行集群式分布式训练的具体功能； + + :param model: 训练时使用的 **PaddlePaddle** 模型； + :param fp16: 是否开启混合精度训练； + :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, fp16: Optional[bool] = False, **kwargs): + def __init__(self, model: "paddle.nn.Layer", fp16: Optional[bool] = False, **kwargs): if not isinstance(model, paddle.nn.Layer): raise ValueError(f"Parameter `model` can not be `{type(model)}` in `PaddleDriver`, it should be exactly " f"`paddle.nn.Layer` type.") @@ -69,10 +84,10 @@ class PaddleDriver(Driver): def zero_grad(self, set_to_none: bool = False): r""" - 实现深度学习中的梯度的置零操作，应当直接通过优化器 optimizers 来将梯度置零； - 注意梯度累积不需要在这里实现，trainer 已经在内部实现了梯度累积； + 实现深度学习中的梯度的置零操作，应当直接通过优化器 ``optimizers`` 来将梯度置零； + 注意梯度累积不需要在这里实现，:class:`~fastNLP.core.Trainer` 已经在内部实现了梯度累积； - :param set_to_none: 用来判断是否需要将梯度直接置为 None；Paddle中这个参数无效。 + :param set_to_none: 用来判断是否需要将梯度直接置为 ``None``；在 **PaddlePaddle** 中这个参数无效。 """ for optimizer in self.optimizers: optimizer.clear_grad() @@ -87,14 +102,6 @@ class PaddleDriver(Driver): @staticmethod def check_dataloader_legality(dataloader, dataloader_name, is_train: bool = False): - r""" - 该函数会在 trainer 或者 evaluator 设置 dataloader 后检测 dataloader 的合法性。 - 要求传入的 dataloader 必须为 `paddle.io.DataLoader` 或包含该类型的字典。 - - :param dataloader: 需要检测的输入的 `dataloader`； - :param dataloader_name: - :param is_train: - """ if is_train: if not isinstance(dataloader, DataLoader): raise ValueError(f"Parameter `{dataloader_name}` should be 'paddle.io.DataLoader' type, not {type(dataloader)}.") @@ -164,16 +171,15 @@ class PaddleDriver(Driver): @rank_zero_call def save_model(self, filepath: str, only_state_dict: bool = True, **kwargs): r""" - 保存模型的函数；注意函数 `save` 是用来进行断点重训的函数； + 将模型保存到 ``filepath`` 中。 :param filepath: 保存文件的文件位置（需要包括文件名）； - :param only_state_dict: 是否只保存模型的 `state_dict`；如果为 False，则会调用 `paddle.jit.save` 函数 - 保存整个模型的参数，此时需要传入 `input_spec` 参数，否则在 load 时会报错。 - :param kwargs: - input_spec: 描述存储模型 forward 方法的输入，当 `only_state_dict` 为 False时必须传入，否则加载时会报错。 - 可以通过 InputSpec 或者示例 Tensor 进行描述。详细的可以参考 paddle 关于`paddle.jit.save` - 的文档： - https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/jit/save_cn.html#save + :param only_state_dict: 是否只保存模型的 ``state_dict``；如果为 ``False``，则会调用 ``paddle.jit.save`` + 函数保存整个模型的参数，此时需要传入 ``input_spec`` 参数； + :kwargs: + * input_spec -- 描述存储模型 ``forward`` 方法的输入； + 当 ``only_state_dict`` 为 ``False`` 时必须传入，否则加载时会报错。您可以通过 ``InputSpec`` 或者示例 ``Tensor`` + 进行描述。详细的使用方法可以参考 **PaddlePaddle** `关于 paddle.jit.save 函数的文档 `_； """ model = self.unwrap_model() if isinstance(filepath, Path): @@ -189,14 +195,6 @@ class PaddleDriver(Driver): paddle.jit.save(model, filepath, input_spec) def load_model(self, filepath: str, only_state_dict: bool = True, **kwargs): - r""" - 加载模型的函数；将 filepath 中的模型加载并赋值给当前 model 。 - - :param filepath: 需要被加载的对象的文件位置（需要包括文件名）； - :param load_state_dict: 保存的文件是否只是模型的权重，还是完整的模型。即便是保存的完整的模型，此处也只能使用尝试加载filepath - 模型中的权重到自身模型，而不会直接替代当前 Driver 中的模型。 - :return: 返回加载指定文件后的结果； - """ model = self.unwrap_model() if isinstance(filepath, Path): filepath = str(filepath) @@ -210,6 +208,27 @@ class PaddleDriver(Driver): @rank_zero_call def save(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() 的时候，需要将 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 函数的文档 `_； + + .. todo: + + 等 Driver 的文档写完 + + """ # 传入的 dataloader 参数是 trainer 的 dataloader 属性，因为 driver 的所有 dataloader 我们是不会去改变它的，而是通过改变 # trainer.dataloader 来改变 dataloader 的状态，从而适配训练或者评测环境； @@ -352,37 +371,41 @@ class PaddleDriver(Driver): r""" 返回一个不计算梯度的环境用来对模型进行评测； - :return: context 上下文对象 `paddle.no_grad`; + :return: 上下文对象 ``paddle.no_grad``; """ return paddle.no_grad @staticmethod def move_model_to_device(model: "paddle.nn.Layer", device: Union[str, int, "paddle.CUDAPlace", "paddle.CPUPlace"]): r""" - 用来将模型转移到指定的 device 上； - 在 Paddle 中使用可能会引起因与设置的设备不一致而产生的问题，请注意。 + 用来将模型 ``model`` 转移到指定的设备上； + + .. note:: + + 在 **Paddle** 中使用可能会引起因与设置的设备不一致而产生的问题，请注意。 + + :param model: 需要进行转移的模型； + :param device: 目标设备； """ if device is not None: model.to(device) - def move_data_to_device(self, batch: "paddle.Tensor"): + def move_data_to_device(self, batch: Any) -> Any: r""" - 将数据迁移到指定的机器上；batch 可能是 list 也可能 dict ，或其嵌套结构。 - 在 Paddle 中使用可能会引起因与设置的设备不一致而产生的问题，请注意。 + 将数据集合 ``batch`` 迁移到指定的机器上。 + + .. note:: + + 在 **Paddle** 中使用可能会引起因与设置的设备不一致而产生的问题，请注意。 - :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) @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 set the seed - with ``seed_everything(seed, workers=True)``. - - See also the PyTorch documentation on - `randomness in DataLoaders `_. - """ # implementation notes: https://github.com/pytorch/pytorch/issues/5059#issuecomment-817392562 global_rank = rank if rank is not None else int(os.environ.get(FASTNLP_GLOBAL_RANK, 0)) # TODO gpu @@ -409,9 +432,6 @@ class PaddleDriver(Driver): @staticmethod def get_dataloader_args(dataloader: "DataLoader"): - """ - 获取 dataloader 的 shuffle 和 drop_last 属性； - """ @dataclass class Res: diff --git a/fastNLP/core/drivers/paddle_driver/single_device.py b/fastNLP/core/drivers/paddle_driver/single_device.py index c0957dbf..10779fd6 100644 --- a/fastNLP/core/drivers/paddle_driver/single_device.py +++ b/fastNLP/core/drivers/paddle_driver/single_device.py @@ -33,9 +33,20 @@ __all__ = [ class PaddleSingleDriver(PaddleDriver): """ - 支持 paddle cpu 或单卡 gpu 训练的 driver + 实现了 **PaddlePaddle** 框架下在单卡或 ``cpu`` 环境下训练功能的 **Driver**。 + + :param model: 训练时使用的 **PaddlePaddle** 模型； + :param device: 训练使用的设备； + :param fp16: 是否开启混合精度训练； + :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: Union[str, int], fp16: Optional[bool] = False, **kwargs): + def __init__(self, model: "paddle.nn.Layer", device: Union[str, int], fp16: Optional[bool] = False, **kwargs): if isinstance(model, DataParallel): raise ValueError("`paddle.DataParallel` is not supported in `PaddleSingleDriver`") @@ -62,7 +73,7 @@ class PaddleSingleDriver(PaddleDriver): def setup(self): r""" - 该函数用来初始化训练环境，用于设置当前训练的设备，并将模型迁移到对应设备上。 + 初始化训练环境；设置当前训练的设备，并将模型迁移到对应设备上。 """ device = _convert_data_device(self.data_device) @@ -127,17 +138,20 @@ class PaddleSingleDriver(PaddleDriver): return dataloader def unwrap_model(self): - if isinstance(self.model, paddle.DataParallel): - return self.model._layers - else: - return self.model + """ + 返回训练使用的模型。 + """ + return self.model @property - def data_device(self): + def data_device(self) -> str: """ - 返回数据所在的设备。由于单卡模式不支持 data_device，因此返回的是 model_device + :return: 数据和模型所在的设备； """ return self.model_device - def is_distributed(self): + def is_distributed(self) -> bool: + """ + 判断是否为分布式的 **Driver** ，在 ``PaddleSingleDriver`` 中，返回 ``False``。 + """ return False diff --git a/fastNLP/core/drivers/paddle_driver/utils.py b/fastNLP/core/drivers/paddle_driver/utils.py index 6362193e..1a324c97 100644 --- a/fastNLP/core/drivers/paddle_driver/utils.py +++ b/fastNLP/core/drivers/paddle_driver/utils.py @@ -31,7 +31,15 @@ __all__ = [ def _select_seed_randomly(min_seed_value: int = 0, max_seed_value: int = 255) -> int: return random.randint(min_seed_value, max_seed_value) -def paddle_seed_everything(seed: Optional[int] = None, workers: bool = False) -> int: +def paddle_seed_everything(seed: Optional[int], workers: bool = False) -> int: + r""" + 为 **paddle**、**numpy**、**python.random** 伪随机数生成器设置种子。 + + :param seed: 全局随机状态的整数值种子。如果为 ``None``，将从环境变量 ``FASTNLP_GLOBAL_SEED`` 中读取种子或随机选择; + :param workers: 如果为 ``True`` ，则会设置环境变量 ``FASTNLP_SEED_WORKERS`` 。该环境变量会在 :class:`~fastNLP.core.Trainer` + 中配置 ``dataloader`` 时用于设置 ``worker_init_fn`` 。如果用户已经为 ``dataloader`` 提供了 ``worker_init_fn`` ，则设置 + 此参数将没有影响; + """ max_seed_value = np.iinfo(np.uint32).max min_seed_value = np.iinfo(np.uint32).min @@ -70,7 +78,7 @@ def paddle_seed_everything(seed: Optional[int] = None, workers: bool = False) -> def reset_seed() -> None: """ - fleet 会开启多个进程，因此当用户在脚本中指定 seed_everything 时，在开启多个脚本后，会在每个脚本内重新 + ``fleet`` 会开启多个进程，因此当用户在脚本中指定 ``seed_everything`` 时，在开启多个脚本后，会在每个脚本内重新 进行随机数的设置； """ seed = os.environ.get(FASTNLP_GLOBAL_SEED, None) @@ -80,8 +88,8 @@ def reset_seed() -> None: class _FleetWrappingModel(Layer): """ - 参考 _DDPWrappingModel ， paddle 的分布式训练也需要用 paddle.nn.DataParallel 进行包装，采用和 - pytorch 相似的处理方式 + 参考 :class:`fastNLP.core.drivers.torch_driver.utils._DDPWrappingModel` ， **PaddlePaddle** 的分布式训练也需要用 :class:`paddle.nn.DataParallel` 进行包装，采用和 + **pytorch** 相似的处理方式 """ def __init__(self, model: 'nn.Layer'): super(_FleetWrappingModel, self).__init__() @@ -100,7 +108,7 @@ class _FleetWrappingModel(Layer): class DummyGradScaler: """ - 用于仿造的GradScaler对象，防止重复写大量的if判断 + 用于仿造的 **GradScaler** 对象，防止重复写大量的if判断 """ def __init__(self, *args, **kwargs): pass @@ -144,7 +152,7 @@ def _build_fp16_env(dummy=False): def find_free_ports(num): """ - 在空闲的端口中找到 num 个端口 + 在空闲的端口中找到 ``num`` 个端口 """ def __free_port(): with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s: @@ -174,8 +182,8 @@ def find_free_ports(num): def replace_batch_sampler(dataloader: "DataLoader", batch_sampler: "BatchSampler"): """ - 利用 `batch_sampler` 重新构建一个 DataLoader，起到替换 `batch_sampler` 又不影响原 `dataloader` 的作用。 - 考虑了用户自己定制了 DataLoader 的情形。 + 利用 ``batch_sampler`` 重新构建一个 ``DataLoader``，起到替换 ``batch_sampler`` 又不影响原 ``dataloader`` 的作用。 + 考虑了用户自己定制了 ``DataLoader`` 的情形。 """ # 拿到非下划线开头的实例属性； instance_attrs = {k: v for k, v in vars(dataloader).items() if not k.startswith('_')} @@ -246,7 +254,7 @@ def replace_batch_sampler(dataloader: "DataLoader", batch_sampler: "BatchSampler def replace_sampler(dataloader, new_sampler): """ - 使用 `new_sampler` 重新构建一个 BatchSampler，并替换到 `dataloader` 中 + 使用 ``new_sampler`` 重新构建一个 ``BatchSampler``，并替换到 ``dataloader`` 中 """ new_batch_sampler = deepcopy(dataloader.batch_sampler) new_batch_sampler.sampler = new_sampler diff --git a/fastNLP/core/drivers/torch_driver/torch_driver.py b/fastNLP/core/drivers/torch_driver/torch_driver.py index 5aee15e9..a1b83d07 100644 --- a/fastNLP/core/drivers/torch_driver/torch_driver.py +++ b/fastNLP/core/drivers/torch_driver/torch_driver.py @@ -297,7 +297,7 @@ class TorchDriver(Driver): sampler = RandomSampler(dataloader_args.sampler.data_source) logger.debug("Replace torch RandomSampler into fastNLP RandomSampler.") elif self.is_distributed(): - raise RuntimeError("It is not allowed to use checkpoint retraining when you do not use our or " + raise RuntimeError("It is not allowed to use checkpoint retraining when you do not use our" "`ReproducibleSampler`.") else: sampler = ReproduceBatchSampler( diff --git a/fastNLP/core/utils/jittor_utils.py b/fastNLP/core/utils/jittor_utils.py index 08b3b7a8..f29b1f46 100644 --- a/fastNLP/core/utils/jittor_utils.py +++ b/fastNLP/core/utils/jittor_utils.py @@ -32,8 +32,8 @@ def is_jittor_dataset(dataset) -> bool: def jittor_collate_wraps(func, auto_collator: Callable): """ - 对 ``jittor`` 的 ``collate_fn`` 进行 ``wrap`` 封装,。如果数据集为 ``mapping`` 类型，那么采用 ``auto_collator`` ，否则 - 还是采用 ``jittor`` 的 ``collate_batch``。 + 对 ``jittor`` 的 ``collate_fn`` 进行 ``wrap`` 封装,。如果数据集为 ``mapping`` 类型，那么采用 ``auto_collator`` ， + 否则还是采用 ``jittor`` 的 ``collate_batch``。 :param func: :param auto_collator: diff --git a/fastNLP/core/utils/paddle_utils.py b/fastNLP/core/utils/paddle_utils.py index d3764d4e..2d7b65cc 100644 --- a/fastNLP/core/utils/paddle_utils.py +++ b/fastNLP/core/utils/paddle_utils.py @@ -61,8 +61,8 @@ def _convert_data_device(device: Union[str, int]) -> str: def paddle_to(data: "paddle.Tensor", device: Union[str, int]) -> "paddle.Tensor": """ - 将 ``data`` 迁移到指定的 ``device`` 上。``paddle.Tensor`` 没有类似 ``torch.Tensor`` 的 ``to`` 函数，该函数 - 只是集成了 :func:`paddle.Tensor.cpu` 和 :func:`paddle.Tensor.cuda` 两个函数。 + 将 ``data`` 迁移到指定的 ``device`` 上。``paddle.Tensor`` 没有类似 ``torch.Tensor`` 的 ``to`` 函数， + 该函数只是集成了 :func:`paddle.Tensor.cpu` 和 :func:`paddle.Tensor.cuda` 两个函数。 :param data: 要迁移的张量； :param device: 目标设备，可以是 ``str`` 或 ``int`` 类型； @@ -130,8 +130,8 @@ def paddle_move_data_to_device(batch: Any, device: Optional[Union[str, int]]) -> 将 **paddle** 的数据集合传输到给定设备。只有 :class:`paddle.Tensor` 对象会被传输到设备中，其余保持不变。 :param batch: 需要进行迁移的数据集合； - :param device: 目标设备。可以是显卡设备的编号，或是``cpu``, ``gpu`` 或 ``gpu:x`` 格式的字符串；当这个参数 - 为 `None`` 时，不会执行任何操作。 + :param device: 目标设备。可以是显卡设备的编号，或是``cpu``, ``gpu`` 或 ``gpu:x`` 格式的字符串； + 当这个参数为 `None`` 时，不会执行任何操作。 :return: 迁移到新设备上的数据集合； """ if device is None: diff --git a/fastNLP/core/utils/rich_progress.py b/fastNLP/core/utils/rich_progress.py index 04cb6383..ff01acbf 100644 --- a/fastNLP/core/utils/rich_progress.py +++ b/fastNLP/core/utils/rich_progress.py @@ -1,6 +1,6 @@ """ -该文件用于为 **fastNLP** 提供一个统一的 ``progress bar`` 管理，通过共用一个``Task`` 对象， :class:`~fastNLP.core.Trainer` 中 -的 ``progress bar`` 和 :class:`~fastNLP.core.Evaluator` 中的 ``progress bar`` 才能不冲突 +该文件用于为 **fastNLP** 提供一个统一的 ``progress bar`` 管理，通过共用一个``Task`` 对象， :class:`~fastNLP.core.Trainer` +中的 ``progress bar`` 和 :class:`~fastNLP.core.Evaluator` 中的 ``progress bar`` 才能不冲突 """ import sys from typing import Any, Union, Optional diff --git a/fastNLP/core/utils/utils.py b/fastNLP/core/utils/utils.py index 4d8bbb5e..5e02ef6d 100644 --- a/fastNLP/core/utils/utils.py +++ b/fastNLP/core/utils/utils.py @@ -60,7 +60,7 @@ def auto_param_call(fn: Callable, *args, signature_fn: Optional[Callable] = None ``value`` 的参数。 1. 该函数用来提供给用户根据字符串匹配从而实现自动调用； - 2. 注意 ``mapping`` 默认为 ``None``，如果你希望指定输入和运行函数的参数的对应方式，那么你应当让 ``mapping`` 为一个字典传入进来； + 2. 注意 ``mapping`` 默认为 ``None``，如果您希望指定输入和运行函数的参数的对应方式，那么您应当让 ``mapping`` 为一个字典传入进来； 如果 ``mapping`` 不为 ``None``，那么我们一定会先使用 ``mapping`` 将输入的字典的 ``keys`` 修改过来，因此请务必亲自检查 ``mapping`` 的正确性； 3. 如果输入的函数的参数有默认值，那么如果之后的输入中没有该参数对应的值，我们就会使用该参数对应的默认值，否则也会使用之后的输入的值； 4. 如果输入的函数是一个 ``partial`` 函数，情况同第三点，即和默认参数的情况相同； @@ -82,8 +82,8 @@ def auto_param_call(fn: Callable, *args, signature_fn: Optional[Callable] = None :param fn: 用来进行实际计算的函数，其参数可以包含有默认值； :param args: 一系列的位置参数，应当为一系列的字典，我们需要从这些输入中提取 ``fn`` 计算所需要的实际参数； - :param signature_fn: 函数，用来替换 ``fn`` 的函数签名，如果该参数不为 ``None``，那么我们首先会从该函数中提取函数签名，然后通过该函数签名提取 - 参数值后，再传给 ``fn`` 进行实际的运算； + :param signature_fn: 函数，用来替换 ``fn`` 的函数签名，如果该参数不为 ``None``，那么我们首先会从该函数中提取函数签名， + 然后通过该函数签名提取参数值后，再传给 ``fn`` 进行实际的运算； :param mapping: 一个字典，用来更改其前面的字典的键值； :return: 返回 ``fn`` 运行的结果； @@ -195,8 +195,8 @@ def _get_fun_msg(fn, with_fp=True)->str: def _check_valid_parameters_number(fn, expected_params:List[str], fn_name=None): """ - 检查一个函数是否需要 expected_params 参数(检测数量是否匹配)。除掉 self （如果是method），给定默认值的参数等。如果匹配不上，就会 - 进行报错。 + 检查一个函数是否需要 expected_params 参数(检测数量是否匹配)。除掉 self （如果是method），给定默认值的参数等。 + 如果匹配不上，就会进行报错。 :param fn: 需要检测的函数，可以是 method 或者 function 。 :param expected_params: 期待应该支持的参数。 @@ -345,8 +345,8 @@ def apply_to_collection( :param dtype: 数据的类型，函数 ``function`` 只会被应用于 ``data`` 中类型为 ``dtype`` 的数据； :param function: 对数据进行处理的函数； :param args: ``function`` 所需要的其它参数； - :param wrong_dtype: ``function`` 一定不会生效的数据类型。如果数据既是 ``wrong_dtype`` 类型又是 ``dtype`` 类型 - 那么也不会生效； + :param wrong_dtype: ``function`` 一定不会生效的数据类型。 + 如果数据既是 ``wrong_dtype`` 类型又是 ``dtype`` 类型那么也不会生效； :param include_none: 是否包含执行结果为 ``None`` 的数据，默认为 ``True``； :param kwargs: ``function`` 所需要的其它参数； :return: 经过 ``function`` 处理后的数据集合； @@ -587,7 +587,7 @@ def seq_len_to_mask(seq_len, max_len: Optional[int]): :param seq_len: 大小为 ``(B,)`` 的长度序列； :param int max_len: 将长度补齐或截断到 ``max_len``。默认情况（为 ``None``）使用的是 ``seq_len`` 中最长的长度； 但在 :class:`torch.nn.DataParallel` 等分布式的场景下可能不同卡的 ``seq_len`` 会有区别，所以需要传入 - 一个 ``max_len`` 使得 ``mask`` 的补齐或截断到该长度。 + ``max_len`` 使得 ``mask`` 的补齐或截断到该长度。 :return: 大小为 ``(B, max_len)`` 的 ``mask``， 元素类型为 ``bool`` 或 ``uint8`` """ if isinstance(seq_len, np.ndarray): diff --git a/fastNLP/modules/mix_modules/utils.py b/fastNLP/modules/mix_modules/utils.py index b19a5d53..142644f9 100644 --- a/fastNLP/modules/mix_modules/utils.py +++ b/fastNLP/modules/mix_modules/utils.py @@ -202,12 +202,12 @@ def jittor2torch(batch: Any, device: str = None, no_gradient: bool = None) -> An .. note:: - 注意，由于 **pytorch** 和 **jittor** 之间的差异，从 :class:`jittor.Var` 转换 - 至 :class:`torch.Tensor` 的过程中无法保留原张量的梯度。 + 注意，由于 **pytorch** 和 **jittor** 之间的差异，从 :class:`jittor.Var` 转换至 + :class:`torch.Tensor` 的过程中无法保留原张量的梯度。 :param batch: 包含 :class:`jittor.Var` 类型的数据集合； :param device: 是否将转换后的张量迁移到特定设备上。为 ``None``时，和输入保持一致； - :param no_gradient: 是否保留原张量的梯度，在这个函数中该参数无效。 + :param no_gradient: 是否保留原张量的梯度，在这个函数中该参数无效; :return: 转换后的数据； """ diff --git a/fastNLP/transformers/torch/file_utils.py b/fastNLP/transformers/torch/file_utils.py index 60f95fdd..4c7ee7a4 100644 --- a/fastNLP/transformers/torch/file_utils.py +++ b/fastNLP/transformers/torch/file_utils.py @@ -82,6 +82,52 @@ def filelock(path): except: pass +class HfFolder: + """ + hugging_face.HfFolder + version = 0.5.1 + """ + path_token = os.path.expanduser("~/.huggingface/token") + + @classmethod + def save_token(cls, token): + """ + Save token, creating folder as needed. + + Args: + token (`str`): + The token to save to the [`HfFolder`] + """ + os.makedirs(os.path.dirname(cls.path_token), exist_ok=True) + with open(cls.path_token, "w+") as f: + f.write(token) + + @classmethod + def get_token(cls): + """ + Retrieves the token + + Returns: + `str` or `None`: The token, `None` if it doesn't exist. + + """ + try: + with open(cls.path_token, "r") as f: + return f.read() + except FileNotFoundError: + pass + + @classmethod + def delete_token(cls): + """ + Deletes the token from storage. Does not fail if token does not exist. + """ + try: + os.remove(cls.path_token) + except FileNotFoundError: + pass + + def is_offline_mode(): return _is_offline_mode @@ -629,11 +675,10 @@ def get_from_cache( if isinstance(use_auth_token, str): headers["authorization"] = f"Bearer {use_auth_token}" elif use_auth_token: - raise RuntimeError("`use_auth_token=True` is not supported in FastNLP now") - # token = HfFolder.get_token() - # if token is None: - # raise EnvironmentError("You specified use_auth_token=True, but a huggingface token was not found.") - # headers["authorization"] = f"Bearer {token}" + token = HfFolder.get_token() + if token is None: + raise EnvironmentError("You specified use_auth_token=True, but a huggingface token was not found.") + headers["authorization"] = f"Bearer {token}" url_to_download = url etag = None @@ -791,13 +836,7 @@ def get_list_of_files( if isinstance(use_auth_token, str): token = use_auth_token elif use_auth_token is True: - # token = HfFolder.get_token() - path_token = os.path.expanduser("~/.huggingface/token") - try: - with open(path_token, "r") as f: - token = f.read() - except FileNotFoundError: - token = None + token = HfFolder.get_token() else: token = None # model_info = HfApi(endpoint=HUGGINGFACE_CO_RESOLVE_ENDPOINT).model_info( diff --git a/fastNLP/transformers/torch/models/auto/tokenization_auto.py b/fastNLP/transformers/torch/models/auto/tokenization_auto.py index e275579f..f1618d6a 100644 --- a/fastNLP/transformers/torch/models/auto/tokenization_auto.py +++ b/fastNLP/transformers/torch/models/auto/tokenization_auto.py @@ -15,7 +15,7 @@ """ Auto Tokenizer class. """ from collections import OrderedDict -from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union +from typing import TYPE_CHECKING, Optional, Tuple from ...file_utils import ( is_sentencepiece_available,