class_docstr: 实现了 **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::
  
      您可以在使用 ``JittorSingleDriver`` 和 ``JittorMPIDriver`` 时使用 ``JittorDriver`` 提供的接口。
  
  :param model: 训练时使用的 **jittor** 模型
  :param fp16: 是否开启混合精度训练
  :param jittor_kwargs:
class_docstr_tok: ["\u5b9e\u73b0\u4e86","*","*","jittor","*","*","\u6846\u67b6\u8bad\u7ec3\u529f\u80fd\u7684\u57fa\u672c","``","Driver","``","\u3002\u8fd9\u4e2a\u7c7b\u88ab\u4ee5\u4e0b\u5b50\u7c7b\u7ee7\u627f\uff1a","1.",":","class",":","`","~fastNLP.core.drivers.jittor_driver.JittorSingleDriver","`","\uff1a\u5b9e\u73b0\u4e86\u4f7f\u7528\u5355\u5361\u548c","``","cpu","``","\u8bad\u7ec3\u7684\u5177\u4f53\u529f\u80fd\uff1b","2.",":","class",":","`","~fastNLP.core.drivers.jittor_driver.JittorMPIDriver","`","\uff1a\u5b9e\u73b0\u4e86\u4f7f\u7528","``","mpi","``","\u542f\u52a8","*","*","jittor","*","*","\u5206\u5e03\u5f0f\u8bad\u7ec3\u7684\u529f\u80fd\uff1b","..","warning",":",":","\u60a8\u4e0d\u5e94\u5f53\u76f4\u63a5\u521d\u59cb\u5316\u8be5\u7c7b\uff0c\u7136\u540e\u4f20\u5165\u7ed9","``","Trainer","``","\uff0c\u6362\u53e5\u8bdd\u8bf4\uff0c\u60a8\u5e94\u5f53\u4f7f\u7528\u8be5\u7c7b\u7684\u5b50\u7c7b","``","JittorSingleDriver","``","\u548c","``","TorchDDPDriver","``","\uff0c\u800c\u4e0d\u662f","\u8be5\u7c7b\u672c\u8eab\u3002","..","note",":",":","\u60a8\u53ef\u4ee5\u5728\u4f7f\u7528","``","JittorSingleDriver","``","\u548c","``","JittorMPIDriver","``","\u65f6\u4f7f\u7528","``","JittorDriver","``","\u63d0\u4f9b\u7684\u63a5\u53e3\u3002",":","param","model",":","\u8bad\u7ec3\u65f6\u4f7f\u7528\u7684","*","*","jittor","*","*","\u6a21\u578b",":","param","fp16",":","\u662f\u5426\u5f00\u542f\u6df7\u5408\u7cbe\u5ea6\u8bad\u7ec3",":","param","jittor_kwargs",":"]
docstr: 断点重训的加载函数，该函数会负责读取数据，并且恢复 **优化器** 、**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)。
docstr_tok: ["\u65ad\u70b9\u91cd\u8bad\u7684\u52a0\u8f7d\u51fd\u6570\uff0c\u8be5\u51fd\u6570\u4f1a\u8d1f\u8d23\u8bfb\u53d6\u6570\u636e\uff0c\u5e76\u4e14\u6062\u590d","*","*","\u4f18\u5316\u5668","*","*","\u3001","*","*","sampler","*","*","\u7684\u72b6\u6001\u548c","*","*","\u6a21\u578b","*","*","\uff08\u5982\u679c","``","should_load_model","``","\u4e3a","True\uff09\u4ee5\u53ca\u5176\u5b83","\u5728",":","meth",":","`","save_checkpoint","`","\u51fd\u6570\u4e2d\u6267\u884c\u7684\u4fdd\u5b58\u64cd\u4f5c\uff0c\u7136\u540e\u5c06\u4e00\u4e2a","state","\u5b57\u5178\u8fd4\u56de\u7ed9",":","class",":","`","~fastNLP.core.controllers.Trainer","`","\uff08","\u5185\u5bb9\u4e3a",":","meth",":","`","save_checkpoint","`","\u63a5\u53d7\u5230\u7684","``","states","``","\uff09\u3002","\u8be5\u51fd\u6570\u5e94\u8be5\u5728\u6240\u6709","rank","\u4e0a\u6267\u884c\u3002",":","param","folder",":","\u8bfb\u53d6\u8be5","folder","\u4e0b\u7684","``","FASTNLP_CHECKPOINT_FILENAME","``","\u6587\u4ef6\u4e0e","``","FASTNLP_MODEL_FILENAME","``","\uff08\u5982\u679c","should_load_model","\u4e3aTrue\uff09\u3002",":","param","dataloader",":","\u5f53\u524d\u7ed9\u5b9a","dataloader\uff0c\u9700\u8981\u6839\u636e\u4fdd\u5b58\u7684","dataloader","\u72b6\u6001\u5408\u7406\u8bbe\u7f6e\u3002\u82e5\u8be5\u503c\u4e3a","``","None","``","\uff0c\u5219\u4e0d\u9700\u8981\u8fd4\u56de","``","'dataloader","'","``","\u4ee5\u53ca","``","'batch_idx_in_epoch","'","``","\u8fd9\u4e24\u4e2a\u503c\u3002",":","param","only_state_dict",":","\u662f\u5426\u4ec5\u8bfb\u53d6\u6a21\u578b\u7684","state_dict","\uff0c\u5f53","``","should_save_model","``","\u4e3a","``","False","``","\uff0c\u8be5\u53c2\u6570\u65e0\u6548\u3002\u5982\u679c\u4e3a","``","True","``","\uff0c\u8bf4\u660e\u4fdd\u5b58\u7684\u5185\u5bb9\u4e3a\u6743\u91cd\uff1b\u5982\u679c\u4e3a","False","\u8bf4\u660e\u4fdd\u5b58\u7684\u662f\u6a21\u578b\uff0c\u4f46\u4e5f\u662f\u901a\u8fc7\u5f53\u524d","Driver","\u7684\u6a21\u578b\u53bb\u52a0\u8f7d\u4fdd\u5b58\u7684\u6a21\u578b\u7684\u6743\u91cd\uff0c\u800c\u4e0d\u662f\u4f7f\u7528\u4fdd\u5b58\u7684\u6a21\u578b\u66ff\u6362\u5f53\u524d\u6a21\u578b\u3002",":","param","should_load_model",":","\u662f\u5426\u5e94\u8be5\u52a0\u8f7d\u6a21\u578b\uff0c\u5982\u679c\u4e3a","``","False","``","\uff0cDriver","\u5c06\u4e0d\u8d1f\u8d23\u52a0\u8f7d\u6a21\u578b\u3002\u82e5\u8be5\u53c2\u6570\u4e3a","``","True","``","\uff0c\u4f46\u5728\u4fdd\u5b58\u7684\u72b6\u6001\u4e2d\u6ca1\u6709","\u627e\u5230\u5bf9\u5e94\u7684\u6a21\u578b\u72b6\u6001\uff0c\u5219\u62a5\u9519\u3002",":","return",":",":","meth",":","`","save_checkpoint","`","\u51fd\u6570\u8f93\u5165\u7684","``","states","``","\u5185\u5bb9\u3002\u9664\u6b64\u4e4b\u5916\uff0c\u8fd8\u8fd4\u56de\u7684\u5185\u5bb9\u6709\uff1a","*","*","dataloader","*","--","\u6839\u636e\u4f20\u5165\u7684","``","dataloader","``","\u4e0e\u8bfb\u53d6\u51fa\u7684\u72b6\u6001\u8bbe\u7f6e\u4e3a\u5408\u7406\u72b6\u6001\u7684","dataloader\u3002\u5728\u5f53\u524d","``","dataloader","``","\u6837\u672c\u6570\u4e0e\u8bfb\u53d6\u51fa\u7684","sampler","\u6837\u672c\u6570","\u4e0d\u4e00\u81f4\u65f6\u62a5\u9519\u3002","*","*","batch_idx_in_epoch","*","--",":","class",":","`","int","`","\u7c7b\u578b\u7684\u6570\u636e\uff0c\u8868\u660e\u5f53\u524d","epoch","\u8fdb\u884c\u5230\u4e86\u7b2c\u51e0\u4e2a","batch","\u3002\u8bf7\u6ce8\u610f\uff0c\u8be5\u503c\u4e0d\u80fd\u4ec5\u901a\u8fc7\u4fdd\u5b58\u7684\u6570\u636e\u4e2d\u8bfb\u53d6\u7684\uff0c\u56e0\u4e3a\u524d\u540e\u4e24\u6b21\u8fd0\u884c\u7684","``","batch_size","``","\u53ef\u80fd\u6709\u53d8\u5316\uff0c\u800c\u5e94\u8be5\u7b26\u5408\u4ee5\u4e0b\u7b49\u5f0f",":",":","\u8fd4\u56de\u7684","dataloader","\u8fd8\u4f1a\u4ea7\u751f\u7684","batch","\u6570\u91cf","+","batch_idx_in_epoch","=","\u539f\u6765\u4e0d\u65ad\u70b9\u8bad\u7ec3\u65f6\u7684","batch","\u7684\u603b\u6570","\u7531\u4e8e","``","\u8fd4\u56de\u7684","dataloader","\u8fd8\u4f1a\u4ea7\u751f\u7684batch\u6570","``","\u5728","``","batch_size","``","\u4e0e","``","drop_last","``","\u53c2\u6570\u7ed9\u5b9a\u7684\u60c5\u51b5\u4e0b\uff0c\u65e0\u6cd5\u6539\u53d8\uff0c\u56e0\u6b64\u53ea\u80fd\u901a\u8fc7\u8c03\u6574","``","batch_idx_in_epoch","``","\u8fd9\u4e2a\u503c\u6765\u4f7f\u7b49\u5f0f\u6210\u7acb\u3002\u4e00\u4e2a\u7b80\u5355\u7684\u8ba1\u7b97\u539f\u5219\u5982\u4e0b\uff1a","*","drop_last","\u4e3a","``","True","``","\u65f6\uff0c\u7b49\u540c\u4e8e","floor","(","sample_in_this_rank\/batch_size",")","-","floor","(","num_left_samples\/batch_size",")","\uff1b","*","drop_last","\u4e3a","``","False","``","\u65f6\uff0c\u7b49\u540c\u4e8e","ceil","(","sample_in_this_rank\/batch_size",")","-","ceil","(","num_left_samples\/batch_size",")","\u3002"]
code: 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 的状态；
      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的状态
  
      # 4. 恢复 sampler 的状态；
      dataloader_args = self.get_dataloader_args(dataloader)
      if dataloader_args.sampler is None:
          sampler = RandomSampler(
              dataloader_args.sampler.dataset,
              shuffle=dataloader_args.shuffle,
          )
      elif isinstance(dataloader_args.sampler, ReproducibleSampler):
          sampler = dataloader_args.sampler
      elif isinstance(dataloader_args.sampler, JittorRandomSampler):
          sampler = RandomSampler(dataloader_args.sampler.dataset)
          logger.debug(
              "Replace jittor RandomSampler into fastNLP RandomSampler."
          )
      elif isinstance(dataloader_args.sampler, JittorSequentialSampler):
          sampler = RandomSampler(
              dataloader_args.sampler.dataset, shuffle=False
          )
          logger.debug(
              "Replace jittor Sampler into fastNLP RandomSampler without shuffle."
          )
      elif self.is_distributed():
          raise RuntimeError(
              "It is not allowed to use checkpoint retraining when you do not use our"
              "`ReproducibleSampler`."
          )
      else:
          raise RuntimeError(
              f"Jittor sampler {type(dataloader_args.sampler)} is not supported now."
          )
      sampler.load_state_dict(states.pop("sampler_states"))
      states["dataloader"] = self.set_dist_repro_dataloader(
          dataloader, sampler
      )
  
      # 4. 修改 trainer_state.batch_idx_in_epoch
      # sampler 是类似 RandomSampler 的sampler，不是 batch_sampler；
      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
  
      states["batch_idx_in_epoch"] = batch_idx_in_epoch
  
      return states
  
imports: import os
  from pathlib import Path
  from typing import Union, Optional, Dict
  from dataclasses import dataclass
  from fastNLP.envs.imports import _NEED_IMPORT_JITTOR
  from fastNLP.core.drivers.driver import Driver
  from fastNLP.core.dataloaders import JittorDataLoader
  from fastNLP.core.dataloaders import OverfitDataLoader
  from fastNLP.core.samplers import ReproducibleSampler, RandomSampler
  from fastNLP.core.log import logger
  from fastNLP.core.utils import apply_to_collection, nullcontext
  from fastNLP.envs import FASTNLP_MODEL_FILENAME, FASTNLP_CHECKPOINT_FILENAME
label_desc: Use image node_id 12 for calling the JittorDriver obj's underlying member method code with example usage: obj.load_checkpoint(folder, dataloader, only_state_dict, should_load_model) and returns: states