githubmirror/ColossalAI

colossalai/interface/model.py

34 строки851 B

Hongxin Liu

[misc] update pre-commit and run all files (#4752)

19 сен 2023, 09:20

19 сен 2023, 09:20079bf3c

import torch.nn as nn


class ModelWrapper(nn.Module):
    """
    A wrapper class to define the common interface used by booster.

    Args:
        module (nn.Module): The model to be wrapped.
    """

    def __init__(self, module: nn.Module) -> None:
        super().__init__()
        self.module = module

    def unwrap(self):
        """
        Unwrap the model to return the original model for checkpoint saving/loading.
        """
        if isinstance(self.module, ModelWrapper):
            return self.module.unwrap()
        return self.module

    def forward(self, *args, **kwargs):
        return self.module(*args, **kwargs)


class AMPModelMixin:
    """This mixin class defines the interface for AMP training."""

    def update_master_params(self):
        """
        Update the master parameters for AMP training.
        """