code: def load(self, program, for_train=True, dtype=None):
      assert os.path.exists(self.checkpoint_dir)
      checkpoint_dir = os.path.abspath(self.checkpoint_dir)
  
      state_dict = {}
      dist_weight_state_dict = {}
      dist_weight_velocity_state_dict = {}
      dist_bias_state_dict = {}
      dist_bias_velocity_state_dict = {}
      for path in os.listdir(checkpoint_dir):
          path = os.path.join(checkpoint_dir, path)
          if not os.path.isfile(path):
              continue
  
          basename = os.path.basename(path)
          name, ext = os.path.splitext(basename)
  
          if ext not in [".pdopt", ".pdparam"]:
              continue
  
          if not for_train and ext == ".pdopt":
              continue
  
          tensor = paddle.load(path, return_numpy=True)
          if dtype:
              assert dtype in ["float32", "float16"]
              tensor = tensor.astype(dtype)
  
          if "dist@" in name and "@rank@" in name:
              if ".w" in name and "velocity" not in name:
                  dist_weight_state_dict[name] = tensor
              elif ".w" in name and "velocity" in name:
                  dist_weight_velocity_state_dict[name] = tensor
              elif ".b" in name and "velocity" not in name:
                  dist_bias_state_dict[name] = tensor
              elif ".b" in name and "velocity" in name:
                  dist_bias_velocity_state_dict[name] = tensor
  
          else:
              state_dict[name] = tensor
  
      if for_train:
          meta_file = os.path.join(checkpoint_dir, "meta.json")
          if not os.path.exists(meta_file):
              logging.error(
                  "Please make sure the checkpoint dir {} exists, and "
                  "parameters in that dir are validating.".format(
                      checkpoint_dir
                  )
              )
              exit()
  
          with open(meta_file, "r") as handle:
              extra_info = json.load(handle)
  
          # Preporcess distributed parameters.
          pretrain_world_size = extra_info["pretrain_world_size"]
          assert pretrain_world_size > 0
          embedding_size = extra_info["embedding_size"]
          assert embedding_size == self.embedding_size
          num_classes = extra_info["num_classes"]
          assert num_classes == self.num_classes
  
          logging.info(
              "Parameters for pre-training: pretrain_world_size ({}), "
              "embedding_size ({}), and num_classes ({}).".format(
                  pretrain_world_size, embedding_size, num_classes
              )
          )
          logging.info(
              "Parameters for inference or fine-tuning: "
              "world_size ({}).".format(self.world_size)
          )
  
          rank_str = "%05d" % self.rank
  
          dist_weight_state_dict = rearrange_weight(
              dist_weight_state_dict,
              pretrain_world_size,
              self.world_size,
          )
          dist_bias_state_dict = rearrange_weight(
              dist_bias_state_dict, pretrain_world_size, self.world_size
          )
          for name, value in dist_weight_state_dict.items():
              if rank_str in name:
                  state_dict[name] = value
          for name, value in dist_bias_state_dict.items():
              if rank_str in name:
                  state_dict[name] = value
  
          if for_train:
              dist_weight_velocity_state_dict = rearrange_weight(
                  dist_weight_velocity_state_dict,
                  pretrain_world_size,
                  self.world_size,
              )
              dist_bias_velocity_state_dict = rearrange_weight(
                  dist_bias_velocity_state_dict,
                  pretrain_world_size,
                  self.world_size,
              )
              for (
                  name,
                  value,
              ) in dist_weight_velocity_state_dict.items():
                  if rank_str in name:
                      state_dict[name] = value
              for name, value in dist_bias_velocity_state_dict.items():
                  if rank_str in name:
                      state_dict[name] = value
  
      program.set_state_dict(state_dict)
      logging.info("Load checkpoint from '{}'. ".format(checkpoint_dir))
      if for_train:
          return extra_info
      else:
          return {}
  
file_path: insightface/recognition/arcface_paddle/static/utils/io.py
imports: import errno
  import os
  import paddle
  import logging
  import numpy
  import shutil
  import json
  from utils.rearrange_weight import rearrange_weight
label_desc: Use image node_id 3 for calling the Checkpoint obj's underlying member method code with example usage: obj.load(program, for_train, dtype) and returns: extra_info, dict