ColossalAI/colossalai/builder/builder.py

#!/usr/bin/env python
# -*- encoding: utf-8 -*-

import inspect
from collections.abc import Iterable

from colossalai.registry import *


def build_from_config(module, config: dict):
    """Returns an object of :class:`module` constructed from `config`.

    :param module: A python or user-defined class
    :type module: class
    :param config: A python dict containing information used in the construction
        of the return object
    :type config: dict
    :raises AssertionError: Raises an AssertionError if `module` is not a class
    :return: An object of :class:`module`
    :rtype: :class:`module`
    """
    assert inspect.isclass(module), 'module must be a class'
    return module(**config)


def build_from_registry(config, registry: Registry):
    """Returns an object constructed from `config`, the type of the object
    is specified by `registry`.

    :param config: A python dict or a :class:`colossalai.context.Config` object 
        containing information used in the construction of the return object
    :type config: dict or :class:`colossalai.context.colossalai.context.Config`
    :param registry: A registry specifying the type of the return object
    :type registry: :class:`Registry`
    :raises AssertionError: Raises an AssertionError if `registry` is not an object
        of :class:`Registry` or `mod_type` in `config` is not found in `registry`
    :raises Exception: Raises an Exception if an error occurred when building
        from registry
    :return: An object specified by `registry`
    :rtype: Python object specified by `registry`
    """
    config_ = config.copy()  # keep the original config untouched
    assert isinstance(
        registry, Registry), f'Expected type Registry but got {type(registry)}'

    mod_type = config_.pop('type')
    assert registry.has(
        mod_type), f'{mod_type} is not found in registry {registry.name}'
    try:
        obj = registry.get_module(mod_type)(**config_)
    except Exception as e:
        print(
            f'An error occurred when building {mod_type} from registry {registry.name}', flush=True)
        raise e

    return obj


def build_layer(config):
    """Returns a layer object of :class:`nn.Module` constructed from `config`.

    :param config: A python dict or a :class:`colossalai.context.Config` object
        containing information used in the construction of the return object
    :type config: dict or :class:`colossalai.context.Config`
    :return: An object of :class:`nn.Module`
    :rtype: :class:`nn.Module`
    """
    return build_from_registry(config, LAYERS)


def build_loss(config):
    """Returns a loss function object of :class:`torch.autograd.Function` constructed 
    from `config`.

    :param config: A python dict or a :class:`colossalai.context.Config` object
        containing information used in the construction of the return object
    :type config: dict or :class:`colossalai.context.Config`
    :return: An object of :class:`torch.autograd.Function`
    :rtype: :class:`torch.autograd.Function`
    """
    return build_from_registry(config, LOSSES)


def build_model(config):
    """Returns a model object of :class:`nn.Module` constructed from `config`.

    :param config: A python dict or a :class:`colossalai.context.Config` object
        containing information used in the construction of the return object
    :type config: dict or :class:`colossalai.context.Config`
    :return: An object of :class:`nn.Module`
    :rtype: :class:`nn.Module`
    """
    return build_from_registry(config, MODELS)


def build_dataset(config):
    """Returns a dataset object of :class:`torch.utils.data.Dataset` constructed 
    from `config`.

    :param config: A python dict or a :class:`colossalai.context.Config` object
        containing information used in the construction of the return object
    :type config: dict or :class:`colossalai.context.Config`
    :return: An object of :class:`torch.utils.data.Dataset`
    :rtype: :class:`torch.utils.data.Dataset`
    """
    return build_from_registry(config, DATASETS)


def build_optimizer(config, model, params: Iterable = None, need_module=False):
    """Returns an optimizer object of :class:`torch.optim.Optimizer` constructed from `config`, 
    'model' and 'params'.

    :param config: A python dict or a :class:`colossalai.context.Config` object 
        containing information used in the construction of the return object
    :type config: dict or :class:`colossalai.context.Config`
    :param model: A model containing parameters for the optimizer 
    :type model: :class:`nn.Module`
    :param params: A dict containing parameters for the optimizer
    :type params: dict, optional
    :param need_module: Indicates whether the optimizer needs a module
    :type params: bool, optional
    :raises AssertionError: Raises an AssertionError if both `model` and `params` are None
    :return: An object of :class:`torch.optim.Optimizer`
    :rtype: :class:`torch.optim.Optimizer`
    """
    assert model is not None or params is not None, 'arguments model and params can not both be None'
    if need_module:
        config['module'] = model
    elif model is not None:
        config['params'] = model.parameters()
    elif params is not None:
        config['params'] = params

    return build_from_registry(config, OPTIMIZERS)


def build_gradient_handler(config, model, optimizer):
    """Returns a gradient handler object of :class:`BaseGradientHandler` constructed from `config`,
    `model` and `optimizer`.

    :param config: A python dict or a :class:`colossalai.context.Config` object
        containing information used in the construction of the return object
    :type config: dict or :class:`colossalai.context.Config`
    :param model: A model containing parameters for the gradient handler
    :type model: :class:`nn.Module`
    :param optimizer: An optimizer object containing parameters for the gradient handler
    :type optimizer: :class:`torch.optim.Optimizer`
    :return: An object of :class:`BaseGradientHandler`
    :rtype: :class:`BaseGradientHandler`
    """
    config_ = config.copy()
    mod_type = config_.pop('type')
    return GRADIENT_HANDLER.get_module(mod_type)(model, optimizer, **config_)


def build_hooks(config, trainer):
    """Returns a hook object of :class:`BaseHook` constructed from `config` and `trainer`.

    :param config: A python dict or a :class:`colossalai.context.Config` object
        containing information used in the construction of the return object
    :type config: dict or :class:`colossalai.context.Config`
    :param trainer: A :class:`Trainer` object containing parameters for the hook
    :type trainer: :class:`Trainer`
    :return: An object of :class:`BaseHook`
    :rtype: :class:`BaseHook`
    """
    config['trainer'] = trainer
    return build_from_registry(config, HOOKS)


def build_transform(config):
    """Returns a transformation object of :class:`torchvision.transforms` constructed
    from `config`.

    :param config: A python dict or a :class:`colossalai.context.Config` object
        containing information used in the construction of the return object
    :type config: dict or :class:`colossalai.context.Config`
    :return: An object of :class:`torchvision.transforms`
    :rtype: :class:`torchvision.transforms`
    """
    return build_from_registry(config, TRANSFORMS)


def build_pipe_alloc_policy(config):
    """Returns a pipeline allocation policy object constructed from `config`.

    :param config: A python dict or a :class:`colossalai.context.Config` object
        containing information used in the construction of the return object
    :type config: dict or :class:`colossalai.context.Config`
    :return: A pipeline allocation policy object
    :rtype: 
    """
    return build_from_registry(config, PIPE_ALLOC_POLICY)


def build_data_sampler(config, dataset):
    """Returns a data sampler object of :class:`colossalai.nn.data.sampler.BaseSampler`
    constructed from `config`.

    :param config: A python dict or a :class:`colossalai.context.Config` object
        containing information used in the construction of the return object
    :type config: dict or :class:`colossalai.context.Config`
    :param dataset: An object of :class:`torch.utils.data.Dataset` containing information
        used in the construction of the return object
    :type dataset: :class:`torch.utils.data.Dataset`
    :return: An object of :class:`colossalai.nn.data.sampler.BaseSampler`
    :rtype: :class:`colossalai.nn.data.sampler.BaseSampler`
    """
    config_ = config.copy()
    mod_type = config_.pop('type')
    return SAMPLERS.get_module(mod_type)(dataset, **config_)


def build_optimizer_wrapper(config, optimizer, model=None):
    """Returns an optimizer wrapper object of :class:`torch.optim.Optimizer` constructed 
    from `config`, `model` and `optimizer`.

    :param config: A python dict or a :class:`colossalai.context.Config` object 
        containing information used in the construction of the return object
    :type config: dict or :class:`colossalai.context.Config`
    :param optimizer: An optimizer object containing parameters for the gradient handler
    :type optimizer: :class:`torch.optim.Optimizer`
    :param model: A model containing parameters for the gradient handler
    :type model: :class:`nn.Module`, optional
    :return: An object of :class:`torch.optim.Optimizer`
    :rtype: :class:`torch.optim.Optimizer`
    """
    config_ = config.copy()
    mod_type = config_.pop('type')

    # LSG: special treatment for zeor level 3
    if mod_type == 'ZeroRedundancyOptimizer_Level_3':
        return OPTIMIZER_WRAPPERS.get_module(mod_type)(model, optimizer, **config_)
    else:
        return OPTIMIZER_WRAPPERS.get_module(mod_type)(optimizer, **config_)


def build_lr_scheduler(config, optimizer, total_steps, num_steps_per_epoch):
    """Returns a learning rate scheduler object of :class:`torch.optim.lr_scheduler` 
    constructed from `config`, `optimizer`, `total_steps` and `num_steps_per_epoch`.

    :param config: A python dict or a :class:`colossalai.context.Config` object 
        containing information used in the construction of the return object
    :type config: dict or :class:`colossalai.context.Config`
    :param optimizer: An optimizer object containing parameters for the learning rate
        scheduler
    :type optimizer: :class:`torch.optim.Optimizer`
    :param total_steps: Number of total steps of the learning rate scheduler
    :type total_steps: int
    :param num_steps_per_epoch: number of steps per epoch of the learning rate scheduler
    :type num_steps_per_epoch: int
    :return: An object of :class:`torch.optim.lr_scheduler`
    :rtype: :class:`torch.optim.lr_scheduler`
    """
    config_ = config.copy()
    mod_type = config_.pop('type')
    # warmup epochs will overwrite warmup steps
    if 'warmup_epochs' in config_:
        warmup_epochs = config_.pop('warmup_epochs')
        config_['warmup_steps'] = int(num_steps_per_epoch * warmup_epochs)
    return LR_SCHEDULERS.get_module(mod_type)(optimizer, total_steps, num_steps_per_epoch=num_steps_per_epoch,
                                              **config_)