fastreid.utils.checkpoint.Checkpointer Class Reference

Inheritance diagram for fastreid.utils.checkpoint.Checkpointer:

Public Member Functions
	__init__ (self, nn.Module model, str save_dir="", *bool save_to_disk=True, **object checkpointables)
	save (self, str name, **dict kwargs)
	load (self, str path)
	has_checkpoint (self)
	get_checkpoint_file (self)
	get_all_checkpoint_files (self)
	resume_or_load (self, str path, *bool resume=True)
	tag_last_checkpoint (self, str last_filename_basename)

Public Attributes
	model
	checkpointables
	logger
	save_dir
	save_to_disk

Protected Member Functions
	_load_file (self, str f)
	_load_model (self, Any checkpoint)
	_convert_ndarray_to_tensor (self, dict state_dict)

Detailed Description

A checkpointer that can save/load model as well as extra checkpointable
objects.

Definition at line 20 of file checkpoint.py.

Constructor & Destructor Documentation

__init__()

fastreid.utils.checkpoint.Checkpointer.__init__	(		self,
		nn.Module	model,
		str	save_dir = "",
		*bool	save_to_disk = True,
		**object	checkpointables )

Args:
    model (nn.Module): model.
    save_dir (str): a directory to save and find checkpoints.
    save_to_disk (bool): if True, save checkpoint to disk, otherwise
        disable saving for this checkpointer.
    checkpointables (object): any checkpointable objects, i.e., objects
        that have the `state_dict()` and `load_state_dict()` method. For
        example, it can be used like
        `Checkpointer(model, "dir", optimizer=optimizer)`.

Definition at line 26 of file checkpoint.py.

    ):
        """
        Args:
            model (nn.Module): model.
            save_dir (str): a directory to save and find checkpoints.
            save_to_disk (bool): if True, save checkpoint to disk, otherwise
                disable saving for this checkpointer.
            checkpointables (object): any checkpointable objects, i.e., objects
                that have the `state_dict()` and `load_state_dict()` method. For
                example, it can be used like
                `Checkpointer(model, "dir", optimizer=optimizer)`.
        """
        if isinstance(model, (DistributedDataParallel, DataParallel)):
            model = model.module
        self.model = model
        self.checkpointables = copy.copy(checkpointables)
        self.logger = logging.getLogger(__name__)
        self.save_dir = save_dir
        self.save_to_disk = save_to_disk
 

Member Function Documentation

_convert_ndarray_to_tensor()

fastreid.utils.checkpoint.Checkpointer._convert_ndarray_to_tensor ( self, dict state_dict )

protected

In-place convert all numpy arrays in the state_dict to torch tensor.
Args:
    state_dict (dict): a state-dict to be loaded to the model.

Definition at line 228 of file checkpoint.py.

    def _convert_ndarray_to_tensor(self, state_dict: dict):
        """
        In-place convert all numpy arrays in the state_dict to torch tensor.
        Args:
            state_dict (dict): a state-dict to be loaded to the model.
        """
        # model could be an OrderedDict with _metadata attribute
        # (as returned by Pytorch's state_dict()). We should preserve these
        # properties.
        for k in list(state_dict.keys()):
            v = state_dict[k]
            if not isinstance(v, np.ndarray) and not isinstance(
                    v, torch.Tensor
            ):
                raise ValueError(
                    "Unsupported type found in checkpoint! {}: {}".format(
                        k, type(v)
                    )
                )
            if not isinstance(v, torch.Tensor):
                state_dict[k] = torch.from_numpy(v)
 
 

_load_file()

fastreid.utils.checkpoint.Checkpointer._load_file ( self, str f )

protected

Load a checkpoint file. Can be overwritten by subclasses to support
different formats.
Args:
    f (str): a locally mounted file path.
Returns:
    dict: with keys "model" and optionally others that are saved by
        the checkpointer dict["model"] must be a dict which maps strings
        to torch.Tensor or numpy arrays.

Definition at line 174 of file checkpoint.py.

    def _load_file(self, f: str):
        """
        Load a checkpoint file. Can be overwritten by subclasses to support
        different formats.
        Args:
            f (str): a locally mounted file path.
        Returns:
            dict: with keys "model" and optionally others that are saved by
                the checkpointer dict["model"] must be a dict which maps strings
                to torch.Tensor or numpy arrays.
        """
        return torch.load(f, map_location=torch.device("cpu"))
 

_load_model()

fastreid.utils.checkpoint.Checkpointer._load_model ( self, Any checkpoint )

protected

Load weights from a checkpoint.
Args:
    checkpoint (Any): checkpoint contains the weights.

Definition at line 187 of file checkpoint.py.

    def _load_model(self, checkpoint: Any):
        """
        Load weights from a checkpoint.
        Args:
            checkpoint (Any): checkpoint contains the weights.
        """
        checkpoint_state_dict = checkpoint.pop("model")
        self._convert_ndarray_to_tensor(checkpoint_state_dict)
 
        # if the state_dict comes from a model that was wrapped in a
        # DataParallel or DistributedDataParallel during serialization,
        # remove the "module" prefix before performing the matching.
        _strip_prefix_if_present(checkpoint_state_dict, "module.")
 
        # work around https://github.com/pytorch/pytorch/issues/24139
        model_state_dict = self.model.state_dict()
        for k in list(checkpoint_state_dict.keys()):
            if k in model_state_dict:
                shape_model = tuple(model_state_dict[k].shape)
                shape_checkpoint = tuple(checkpoint_state_dict[k].shape)
                if shape_model != shape_checkpoint:
                    self.logger.warning(
                        "'{}' has shape {} in the checkpoint but {} in the "
                        "model! Skipped.".format(
                            k, shape_checkpoint, shape_model
                        )
                    )
                    checkpoint_state_dict.pop(k)
 
        incompatible = self.model.load_state_dict(
            checkpoint_state_dict, strict=False
        )
        if incompatible.missing_keys:
            self.logger.info(
                get_missing_parameters_message(incompatible.missing_keys)
            )
        if incompatible.unexpected_keys:
            self.logger.info(
                get_unexpected_parameters_message(incompatible.unexpected_keys)
            )
 

get_all_checkpoint_files()

fastreid.utils.checkpoint.Checkpointer.get_all_checkpoint_files

(

self

)

Returns:
    list: All available checkpoint files (.pth files) in target
        directory.

Definition at line 135 of file checkpoint.py.

    def get_all_checkpoint_files(self):
        """
        Returns:
            list: All available checkpoint files (.pth files) in target
                directory.
        """
        all_model_checkpoints = [
            os.path.join(self.save_dir, file)
            for file in PathManager.ls(self.save_dir)
            if PathManager.isfile(os.path.join(self.save_dir, file))
               and file.endswith(".pth")
        ]
        return all_model_checkpoints
 

get_checkpoint_file()

fastreid.utils.checkpoint.Checkpointer.get_checkpoint_file

(

self

)

Returns:
    str: The latest checkpoint file in target directory.

Definition at line 120 of file checkpoint.py.

    def get_checkpoint_file(self):
        """
        Returns:
            str: The latest checkpoint file in target directory.
        """
        save_file = os.path.join(self.save_dir, "last_checkpoint")
        try:
            with PathManager.open(save_file, "r") as f:
                last_saved = f.read().strip()
        except IOError:
            # if file doesn't exist, maybe because it has just been
            # deleted by a separate process
            return ""
        return os.path.join(self.save_dir, last_saved)
 

has_checkpoint()

fastreid.utils.checkpoint.Checkpointer.has_checkpoint

(

self

)

Returns:
    bool: whether a checkpoint exists in the target directory.

Definition at line 112 of file checkpoint.py.

    def has_checkpoint(self):
        """
        Returns:
            bool: whether a checkpoint exists in the target directory.
        """
        save_file = os.path.join(self.save_dir, "last_checkpoint")
        return PathManager.exists(save_file)
 

load()

fastreid.utils.checkpoint.Checkpointer.load	(		self,
		str	path )

Load from the given checkpoint. When path points to network file, this
function has to be called on all ranks.
Args:
    path (str): path or url to the checkpoint. If empty, will not load
        anything.
Returns:
    dict:
        extra data loaded from the checkpoint that has not been
        processed. For example, those saved with
        :meth:`.save(**extra_data)`.

Definition at line 77 of file checkpoint.py.

    def load(self, path: str):
        """
        Load from the given checkpoint. When path points to network file, this
        function has to be called on all ranks.
        Args:
            path (str): path or url to the checkpoint. If empty, will not load
                anything.
        Returns:
            dict:
                extra data loaded from the checkpoint that has not been
                processed. For example, those saved with
                :meth:`.save(**extra_data)`.
        """
        print(f'path = {path}')
        if not path:
            # no checkpoint provided
            self.logger.info(
                "No checkpoint found. Training model from scratch"
            )
            return {}
        self.logger.info("Loading checkpoint from {}".format(path))
        if not os.path.isfile(path):
            path = PathManager.get_local_path(path)
            assert os.path.isfile(path), "Checkpoint {} not found!".format(path)
 
        checkpoint = self._load_file(path)
        self._load_model(checkpoint)
        for key, obj in self.checkpointables.items():
            if key in checkpoint:
                self.logger.info("Loading {} from {}".format(key, path))
                obj.load_state_dict(checkpoint.pop(key))
 
        # return any further checkpoint data
        return checkpoint
 

resume_or_load()

fastreid.utils.checkpoint.Checkpointer.resume_or_load	(		self,
		str	path,
		*bool	resume = True )

If `resume` is True, this method attempts to resume from the last
checkpoint, if exists. Otherwise, load checkpoint from the given path.
This is useful when restarting an interrupted training job.
Args:
    path (str): path to the checkpoint.
    resume (bool): if True, resume from the last checkpoint if it exists.
Returns:
    same as :meth:`load`.

Definition at line 149 of file checkpoint.py.

    def resume_or_load(self, path: str, *, resume: bool = True):
        """
        If `resume` is True, this method attempts to resume from the last
        checkpoint, if exists. Otherwise, load checkpoint from the given path.
        This is useful when restarting an interrupted training job.
        Args:
            path (str): path to the checkpoint.
            resume (bool): if True, resume from the last checkpoint if it exists.
        Returns:
            same as :meth:`load`.
        """
        if resume and self.has_checkpoint():
            path = self.get_checkpoint_file()
        return self.load(path)
 

save()

fastreid.utils.checkpoint.Checkpointer.save	(		self,
		str	name,
		**dict	kwargs )

Dump model and checkpointables to a file.
Args:
    name (str): name of the file.
    kwargs (dict): extra arbitrary data to save.

Definition at line 53 of file checkpoint.py.

    def save(self, name: str, **kwargs: dict):
        """
        Dump model and checkpointables to a file.
        Args:
            name (str): name of the file.
            kwargs (dict): extra arbitrary data to save.
        """
        if not self.save_dir or not self.save_to_disk:
            return
 
        data = {}
        data["model"] = self.model.state_dict()
        for key, obj in self.checkpointables.items():
            data[key] = obj.state_dict()
        data.update(kwargs)
 
        basename = "{}.pth".format(name)
        save_file = os.path.join(self.save_dir, basename)
        assert os.path.basename(save_file) == basename, basename
        self.logger.info("Saving checkpoint to {}".format(save_file))
        with PathManager.open(save_file, "wb") as f:
            torch.save(data, f)
        self.tag_last_checkpoint(basename)
 

tag_last_checkpoint()

fastreid.utils.checkpoint.Checkpointer.tag_last_checkpoint	(		self,
		str	last_filename_basename )

Tag the last checkpoint.
Args:
    last_filename_basename (str): the basename of the last filename.

Definition at line 164 of file checkpoint.py.

    def tag_last_checkpoint(self, last_filename_basename: str):
        """
        Tag the last checkpoint.
        Args:
            last_filename_basename (str): the basename of the last filename.
        """
        save_file = os.path.join(self.save_dir, "last_checkpoint")
        with PathManager.open(save_file, "w") as f:
            f.write(last_filename_basename)
 

Member Data Documentation

checkpointables

fastreid.utils.checkpoint.Checkpointer.checkpointables

Definition at line 48 of file checkpoint.py.

logger

fastreid.utils.checkpoint.Checkpointer.logger

Definition at line 49 of file checkpoint.py.

model

fastreid.utils.checkpoint.Checkpointer.model

Definition at line 47 of file checkpoint.py.

save_dir

fastreid.utils.checkpoint.Checkpointer.save_dir

Definition at line 50 of file checkpoint.py.

save_to_disk

fastreid.utils.checkpoint.Checkpointer.save_to_disk

Definition at line 51 of file checkpoint.py.

---

The documentation for this class was generated from the following file:

• smreid/fastreid/utils/checkpoint.py