| |
|
|
| 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) |
|
|
