def compute_module_sizes(
        model: nn.Module,
        dtype: Union[str, torch.dtype, None] = None,
        special_dtypes: Optional[Dict[str, Union[str, torch.dtype]]] = None):
    """Compute the size of each submodule of a given model."""

    def get_dtype(dtype):
        if isinstance(dtype, str):
            dtype = getattr(torch, dtype)
        if dtype is not None:
            assert issubclass(dtype, torch.dtype)
        return dtype

    def dtype_bytes(dtype: torch.dtype):
        if dtype is torch.bool:
            return 1
        if dtype.is_floating_point:
            return torch.finfo(dtype).bits / 8
        else:
            return torch.iinfo(dtype).bits / 8

    if dtype is not None:
        dtype = get_dtype(dtype)
        dtype_size = dtype_bytes(dtype)

    if special_dtypes is not None:
        special_dtypes = {
            key: dtype_bytes(dtype)
            for key, dtype in special_dtypes.items()
        }

    module_sizes = defaultdict(int)
    for name, tensor in chain(
            model.named_parameters(recurse=True),
            model.named_buffers(recurse=True)):
        if special_dtypes is not None and name in special_dtypes:
            size = tensor.numel() * special_dtypes[name]
        elif dtype is None:
            size = tensor.numel() * tensor.element_size()
        else:
            size = tensor.numel() * min(dtype_size, tensor.element_size())
        name_parts = name.split('.')
        for idx in range(len(name_parts) + 1):
            module_sizes['.'.join(name_parts[:idx])] += size

    return module_sizes
Compute the size of each submodule of a given model.
mmpretrain/apis/utils.py