class_docstr: 测试多卡情况下 save 和 load 相关函数的表现
class_docstr_tok: ["\u6d4b\u8bd5\u591a\u5361\u60c5\u51b5\u4e0b","save","\u548c","load","\u76f8\u5173\u51fd\u6570\u7684\u8868\u73b0"]
docstr: 测试 save_model 和 load_model 函数
docstr_tok: ["\u6d4b\u8bd5","save_model","\u548c","load_model","\u51fd\u6570"]
code: def test_save_and_load_model(self, only_state_dict):
      """
      测试 save_model 和 load_model 函数
      """
      try:
          path = "model"
  
          dataloader = DataLoader(self.dataset, batch_size=2)
          self.driver1, self.driver2 = generate_driver(
              40, 1
          ), generate_driver(40, 1)
  
          if only_state_dict:
              self.driver1.save_model(path, only_state_dict)
          else:
              self.driver1.save_model(
                  path,
                  only_state_dict,
                  input_spec=[paddle.ones((4, 1))],
              )
  
          # 同步
          dist.barrier()
          self.driver2.load_model(path, only_state_dict)
  
          for idx, batch in enumerate(dataloader):
              batch = self.driver1.move_data_to_device(batch)
              res1 = self.driver1.model(
                  batch,
                  fastnlp_fn=self.driver1.model._layers.model.evaluate_step,
                  # Driver.model -> DataParallel._layers -> _FleetWrappingModel.model
                  fastnlp_signature_fn=None,
                  wo_auto_param_call=False,
              )
              res2 = self.driver2.model(
                  batch,
                  fastnlp_fn=self.driver2.model._layers.model.evaluate_step,
                  fastnlp_signature_fn=None,
                  wo_auto_param_call=False,
              )
  
              assert paddle.equal_all(res1["pred"], res2["pred"])
      finally:
          if only_state_dict:
              rank_zero_rm(path)
          else:
              rank_zero_rm(path + ".pdiparams")
              rank_zero_rm(path + ".pdiparams.info")
              rank_zero_rm(path + ".pdmodel")
  
imports: import pytest
  from pathlib import Path
  from fastNLP.core.drivers.paddle_driver.fleet import PaddleFleetDriver
  from fastNLP.core.samplers import RandomSampler, UnrepeatedSampler, BucketedBatchSampler, UnrepeatedRandomSampler, UnrepeatedSequentialSampler
  from tests.helpers.models.paddle_model import PaddleNormalModel_Classification_1
  from tests.helpers.datasets.paddle_data import PaddleNormalDataset, PaddleNormalXYDataset
  from tests.helpers.utils import magic_argv_env_context, recover_logger
  from fastNLP.envs.distributed import rank_zero_rm
  from fastNLP import prepare_paddle_dataloader
  from fastNLP.core.drivers.paddle_driver.dist_utils import fastnlp_paddle_all_gather
  from fastNLP.envs.imports import _NEED_IMPORT_PADDLE
  from fastNLP import logger
label_desc: Use image node_id 3 for calling the TestSaveLoad obj's underlying member method code with example usage: obj.test_save_and_load_model(only_state_dict) without return types