| | |
| |
|
| | import logging |
| | import re |
| | from functools import wraps |
| | from re import Pattern |
| | from typing import Callable, Dict, Optional, Tuple |
| |
|
| | import numpy as np |
| | import torch as T |
| |
|
| | from .modules import TTCompressedLinear |
| |
|
| |
|
| | def map_module(root: T.nn.Module, |
| | func: Callable[[T.nn.Module, str], T.nn.Module], |
| | patt: Optional[str] = None) -> T.nn.Module: |
| | """Function ``map_module`` applies a function to each leaf of module tree |
| | which matches to a specified pattern. |
| | |
| | Parameters |
| | ---------- |
| | root : torch.nn.Module |
| | Module to modify. |
| | func : callable |
| | Function to be applied to every module (or matched to pattern) in |
| | module tree. |
| | patt : str, optional |
| | Pattern to filter modules by path in module tree. |
| | |
| | Returns |
| | ------- |
| | torch.nn.Module |
| | Module modified in-place. |
| | """ |
| | @wraps(func) |
| | def func_safe(*args, **kwargs): |
| | node = func(*args, **kwargs) |
| | if not isinstance(node, T.nn.Module): |
| | raise ValueError('Mapped result must be toch.nn.Module type ' |
| | f'but given {type(node)}.') |
| | return node |
| |
|
| | return _map_module(root, func_safe, re.compile(patt or r'.*'), '') |
| |
|
| |
|
| | def _map_module(root: T.nn.Module, |
| | func: Callable[[T.nn.Module, str], T.nn.Module], patt: Pattern, |
| | path: str) -> T.nn.Module: |
| | for name, child in root.named_children(): |
| | node = _map_module(child, func, patt, f'{path}/{name}') |
| | if node != child: |
| | setattr(root, name, node) |
| | if patt.match(path or '/'): |
| | root = func(root, path or '/') |
| | return root |
| |
|
| |
|
| | def convert_linear(module: T.nn.Linear, ctor, **kwargs) -> T.nn.Module: |
| | """Function convert_linear takes module and returns linear module with |
| | approximate matmul. Non-linear modules are returned intact. |
| | """ |
| | if not isinstance(module, T.nn.Linear): |
| | return module |
| | raise NotImplementedError |
| |
|
| |
|
| | def numel(module: T.nn.Module): |
| | value = sum(x.numel() for x in module.parameters()) + \ |
| | sum(x.numel() for x in module.buffers()) |
| |
|
| | def account_prunned(module: T.nn.Module, path: str): |
| | nonlocal value |
| | for name, attr in vars(module).items(): |
| | if not name.endswith('_mask') or not isinstance(attr, T.Tensor): |
| | continue |
| |
|
| | weight_name = name[:-5] |
| | if not hasattr(module, weight_name): |
| | continue |
| |
|
| | weight = getattr(module, weight_name) |
| | value -= weight.numel() - attr.sum() |
| | value += attr.numel() |
| | return module |
| |
|
| | def account_quantized(module: T.nn.Module, path: str): |
| | nonlocal value |
| | if isinstance(module, T.nn.quantized.Linear): |
| | value += module.weight().numel() |
| | if module.bias() is not None: |
| | value += module.bias().numel() |
| | return module |
| |
|
| | def account_rest(module: T.nn.Module, path: str): |
| | account_prunned(module, path) |
| | account_quantized(module, path) |
| | return module |
| |
|
| | map_module(module, account_rest) |
| | return value |
| |
|
| |
|
| | def sizeof(module: T.nn.Module): |
| | value = sum(x.numel() * x.element_size() for x in module.parameters()) + \ |
| | sum(x.numel() * x.element_size() for x in module.buffers()) |
| |
|
| | def account_prunned(module: T.nn.Module, path: str): |
| | nonlocal value |
| | for name, attr in vars(module).items(): |
| | if not name.endswith('_mask') or not isinstance(attr, T.Tensor): |
| | continue |
| |
|
| | weight_name = name[:-5] |
| | if not hasattr(module, weight_name): |
| | continue |
| |
|
| | weight = getattr(module, weight_name) |
| | value -= (weight.numel() - attr.sum()) * weight.element_size() |
| | value += attr.numel() * attr.element_size() |
| | return module |
| |
|
| | def account_quantized(module: T.nn.Module, path: str): |
| | nonlocal value |
| | if isinstance(module, T.nn.quantized.Linear): |
| | value += module.weight().numel() * module.weight().element_size() |
| | if (bias := module.bias()) is not None: |
| | value += bias.numel() * bias.element_size() |
| | return module |
| |
|
| | def account_rest(module: T.nn.Module, path: str): |
| | account_prunned(module, path) |
| | account_quantized(module, path) |
| | return module |
| |
|
| | map_module(module, account_rest) |
| | return value |
| |
|
| |
|
| | def flatten_module(module: T.nn.Module, regexp=None) -> Dict[str, T.nn.Module]: |
| | modules = {} |
| | map_module(module, lambda x, y: modules.update(**{y: x}) or x, regexp) |
| | return modules |
| |
|
| |
|
| | def print_flatten(module: T.nn.Module): |
| | paths = [] |
| | path_len = 0 |
| | names = [] |
| | name_len = 0 |
| | indx_len = 0 |
| |
|
| | def func(module, path): |
| | nonlocal path_len, name_len, indx_len |
| | paths.append(path) |
| | path_len = max(path_len, len(path)) |
| | name = module.__class__.__name__ |
| | names.append(name) |
| | name_len = max(name_len, len(name)) |
| | indx_len += 1 |
| | return module |
| |
|
| | map_module(module, func) |
| |
|
| | indx_len = int(np.ceil(np.log10(indx_len))) |
| | fmt = f'{{indx:>{indx_len}s}} {{path:{path_len}s}} {{name:{name_len}s}}' |
| | print(fmt.format(indx='#', path='Path', name='Layer')) |
| | print('-' * (indx_len + path_len + name_len + 2)) |
| | for i, (path, name) in enumerate(zip(paths, names)): |
| | print(fmt.format(indx=str(i), path=path, name=name)) |
| |
|
| |
|
| | def compress_linear_tt(module: T.nn.Module, path: str, |
| | shape: Tuple[Tuple[int], Tuple[int]], |
| | rank: int) -> T.nn.Module: |
| | if not isinstance(module, T.nn.Linear): |
| | return module |
| |
|
| | |
| | inp_size = np.prod(shape[0]) |
| | out_size = np.prod(shape[1]) |
| | if inp_size == module.in_features and out_size == module.out_features: |
| | pass |
| | elif inp_size == module.out_features and out_size == module.in_features: |
| | shape = (shape[1], shape[0]) |
| | else: |
| | raise ValueError( |
| | 'Input and output features does not match to compression shape: ' |
| | f'{shape[0]} vs {module.in_features} and {shape[1]} vs ' |
| | f'{module.out_features}.') |
| |
|
| | logging.info('apply tt compression to layer %s', path) |
| | return TTCompressedLinear.from_linear(module, shape, rank) |
| |
|
