class_docstr: 实现了 **PaddlePaddle** 框架训练功能的基本 Driver。
  
  这个类被以下子类继承：
  
  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`` 提供的接口。
  
  :param model: 训练时使用的 **PaddlePaddle** 模型
  :param fp16: 是否开启混合精度训练
  :param paddle_kwargs:
docstr: 断点重训的加载函数，该函数会负责读取数据，并且恢复 **优化器** 、**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)。
code: 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 的状态；
      optimizers_state_dict = states.pop("optimizers_state_dict")
      self.load_optimizer_state(optimizers_state_dict)
  
      # 2. 加载模型状态；
      if should_load_model:
          self.load_model(
              folder.joinpath(FASTNLP_MODEL_FILENAME), 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 的状态；
      dataloader_args = self.get_dataloader_args(dataloader)
      if isinstance(
          dataloader_args.batch_sampler, ReproducibleBatchSampler
      ):
          sampler = dataloader_args.batch_sampler
      elif isinstance(dataloader_args.sampler, ReproducibleSampler):
          sampler = dataloader_args.sampler
      elif isinstance(dataloader_args.sampler, PaddleRandomSampler):
          sampler = RandomSampler(dataloader_args.sampler.data_source)
          logger.debug(
              "Replace paddle 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 "
              "`ReproducibleSampler`."
          )
      else:
          sampler = ReproduceBatchSampler(
              batch_sampler=dataloader_args.batch_sampler
              if dataloader_args.batch_sampler is not None
              else dataloader_args.sampler,
              batch_size=dataloader_args.batch_size,
              drop_last=dataloader_args.drop_last,
          )
      sampler.load_state_dict(states.pop("sampler_states"))
      states["dataloader"] = self.set_dist_repro_dataloader(
          dataloader, sampler
      )
  
      # 5. 修改 trainer_state.batch_idx_in_epoch
      # sampler 是类似 RandomSampler 的sampler，不是 batch_sampler；
      if not isinstance(sampler, ReproducibleBatchSampler):
          if dataloader_args.drop_last:
              batch_idx_in_epoch = (
                  len(sampler) // dataloader_args.batch_size
                  - sampler.num_left_samples
                  // dataloader_args.batch_size
              )
          else:
              batch_idx_in_epoch = (
                  len(sampler) + dataloader_args.batch_size - 1
              ) // dataloader_args.batch_size - (
                  sampler.num_left_samples
                  + dataloader_args.batch_size
                  - 1
              ) // dataloader_args.batch_size
      # sampler 是 batch_sampler；
      else:
          batch_idx_in_epoch = sampler.batch_idx_in_epoch
  
      states["batch_idx_in_epoch"] = batch_idx_in_epoch
  
      return states
  
imports: import os
  import random
  from typing import Union, Optional, Dict, Any
  from pathlib import Path
  from functools import partial
  from dataclasses import dataclass
  import numpy
  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
  from fastNLP.core.utils import apply_to_collection, paddle_move_data_to_device
  from fastNLP.core.utils.paddle_utils import _convert_data_device
  from fastNLP.envs import FASTNLP_MODEL_FILENAME, FASTNLP_CHECKPOINT_FILENAME, FASTNLP_GLOBAL_RANK, rank_zero_call
  from fastNLP.core.log import logger
  from fastNLP.core.dataloaders import OverfitDataLoader
  from fastNLP.core.samplers import ReproducibleBatchSampler, ReproducibleSampler, ReproduceBatchSampler, RandomSampler
label_desc: Use image node_id 14 for calling the PaddleDriver obj's underlying member method code with example usage: obj.load_checkpoint(folder, dataloader, only_state_dict, should_load_model) and returns: states