Repository: microsoft/nni
File: nni/contrib/compression/base/wrapper.py (Python)

Code before:

from .apply_method import pruning_apply_methods, quant_apply_methods
from .config import select_modules_by_config
from .setting import INPUT_PREFIX, OUTPUT_PREFIX, canonicalize_settings
from .target_space import (
    TargetSpace,
    TargetType,
    PruningTargetSpace,
    QuantizationTargetSpace,
    DistillationTargetSpace
)
from .fuse_modules import fuse_modules
from ..utils import (
    check_bias,
    update_config,
    validate_fused_modules_config,
    get_fused_module_list,
    get_nested_attr
)

_logger = logging.getLogger(__name__)
OUTPUT_FORMAT = Union[Tensor, Any, Tuple[Tensor, Any], Dict[str, Union[Tensor, Any]]]


class ModuleWrapper(torch.nn.Module):
    def __init__(self, module: torch.nn.Module, module_name: str, config: Dict[str, List[Dict[str, Any]]] | None = None, \
                 fused_modules: List[torch.nn.Module] | None = None):
        """
        Two changes will be done during initialization. One is an attribute named ``_nni_wrapper`` will be set to original module,
        this attribute points to this wrapper in the original module.
        The other is the original ``module.forward`` will be replaced by ``module._nni_wrapper.forward``.

        The module can be unwrapped by ``module._nni_wrapper.unwrap()``.

        Parameters
        ----------
        module
            The torch.nn.Module to be wrapped.
        module_name
            The name of the module in the original model.
        config
            The config is a dict which contains keys (not required): ``pruning``, ``quantization``, ``distillation``.
        fused_modules:
            The List contains a series of modules which need to fuse.
        """
        super().__init__()

        # origin layer information
        assert isinstance(module, torch.nn.Module)
        object.__setattr__(self, 'module', module)
        self.module: torch.nn.Module

Review comment (suggestion):
I think we don't need to modify the api here, just modify `ModuleWrapper.config` -> `ModuleWrapper.configs` to minimize the risk.

Code after:

from .apply_method import pruning_apply_methods, quant_apply_methods
from .config import select_modules_by_config
from .setting import INPUT_PREFIX, OUTPUT_PREFIX, canonicalize_settings
from .target_space import (
    TargetSpace,
    TargetType,
    PruningTargetSpace,
    QuantizationTargetSpace,
    DistillationTargetSpace
)
from .fuse_modules import fuse_modules
from ..utils import (
    check_bias,
    update_config,
    validate_fused_modules_config,
    get_fused_module_list,
    get_nested_attr
)

_logger = logging.getLogger(__name__)
OUTPUT_FORMAT = Union[Tensor, Any, Tuple[Tensor, Any], Dict[str, Union[Tensor, Any]]]


class ModuleWrapper(torch.nn.Module):
    def __init__(self, module: torch.nn.Module, module_name: str, config: Dict[str, Dict[str, Any]] | None = None, \
                 fused_modules: List[torch.nn.Module] | None = None):
        """
        Two changes will be done during initialization. One is an attribute named ``_nni_wrapper`` will be set to original module,
        this attribute points to this wrapper in the original module.
        The other is the original ``module.forward`` will be replaced by ``module._nni_wrapper.forward``.

        The module can be unwrapped by ``module._nni_wrapper.unwrap()``.

        Parameters
        ----------
        module
            The torch.nn.Module to be wrapped.
        module_name
            The name of the module in the original model.
        config
            The config is a dict which contains keys (not required): ``pruning``, ``quantization``, ``distillation``.
        fused_modules:
            The List contains a series of modules which need to fuse.
        """
        super().__init__()

        # origin layer information
        assert isinstance(module, torch.nn.Module)
        object.__setattr__(self, 'module', module)
        self.module: torch.nn.Module