183 lines
6.8 KiB
Python
183 lines
6.8 KiB
Python
|
import torch
|
||
|
|
import copy
|
||
|
|
from typing import Dict, Any
|
||
|
|
|
||
|
|
__all__ = [
|
||
|
|
"set_module_weight",
|
||
|
|
"set_module_bias",
|
||
|
|
"get_module_weight",
|
||
|
|
"get_module_bias",
|
||
|
|
"max_over_ndim",
|
||
|
|
"min_over_ndim",
|
||
|
|
"channel_range",
|
||
|
|
"cross_layer_equalization",
|
||
|
|
"equalize",
|
||
|
|
"converged",
|
||
|
|
]
|
||
|
|
|
||
|
|
_supported_types = {torch.nn.Conv2d, torch.nn.Linear}
|
||
|
|
_supported_intrinsic_types = {torch.ao.nn.intrinsic.ConvReLU2d, torch.ao.nn.intrinsic.LinearReLU}
|
||
|
|
_all_supported_types = _supported_types.union(_supported_intrinsic_types)
|
||
|
|
|
||
|
|
def set_module_weight(module, weight) -> None:
|
||
|
|
if type(module) in _supported_types:
|
||
|
|
module.weight = torch.nn.Parameter(weight)
|
||
|
|
else:
|
||
|
|
module[0].weight = torch.nn.Parameter(weight)
|
||
|
|
|
||
|
|
def set_module_bias(module, bias) -> None:
|
||
|
|
if type(module) in _supported_types:
|
||
|
|
module.bias = torch.nn.Parameter(bias)
|
||
|
|
else:
|
||
|
|
module[0].bias = torch.nn.Parameter(bias)
|
||
|
|
|
||
|
|
def get_module_weight(module):
|
||
|
|
if type(module) in _supported_types:
|
||
|
|
return module.weight
|
||
|
|
else:
|
||
|
|
return module[0].weight
|
||
|
|
|
||
|
|
def get_module_bias(module):
|
||
|
|
if type(module) in _supported_types:
|
||
|
|
return module.bias
|
||
|
|
else:
|
||
|
|
return module[0].bias
|
||
|
|
|
||
|
|
def max_over_ndim(input, axis_list, keepdim=False):
|
||
|
|
"""Apply 'torch.max' over the given axes."""
|
||
|
|
axis_list.sort(reverse=True)
|
||
|
|
for axis in axis_list:
|
||
|
|
input, _ = input.max(axis, keepdim)
|
||
|
|
return input
|
||
|
|
|
||
|
|
def min_over_ndim(input, axis_list, keepdim=False):
|
||
|
|
"""Apply 'torch.min' over the given axes."""
|
||
|
|
axis_list.sort(reverse=True)
|
||
|
|
for axis in axis_list:
|
||
|
|
input, _ = input.min(axis, keepdim)
|
||
|
|
return input
|
||
|
|
|
||
|
|
def channel_range(input, axis=0):
|
||
|
|
"""Find the range of weights associated with a specific channel."""
|
||
|
|
size_of_tensor_dim = input.ndim
|
||
|
|
axis_list = list(range(size_of_tensor_dim))
|
||
|
|
axis_list.remove(axis)
|
||
|
|
|
||
|
|
mins = min_over_ndim(input, axis_list)
|
||
|
|
maxs = max_over_ndim(input, axis_list)
|
||
|
|
|
||
|
|
assert mins.size(0) == input.size(axis), "Dimensions of resultant channel range does not match size of requested axis"
|
||
|
|
return maxs - mins
|
||
|
|
|
||
|
|
def cross_layer_equalization(module1, module2, output_axis=0, input_axis=1):
|
||
|
|
"""Scale the range of Tensor1.output to equal Tensor2.input.
|
||
|
|
|
||
|
|
Given two adjacent tensors', the weights are scaled such that
|
||
|
|
the ranges of the first tensors' output channel are equal to the
|
||
|
|
ranges of the second tensors' input channel
|
||
|
|
"""
|
||
|
|
if type(module1) not in _all_supported_types or type(module2) not in _all_supported_types:
|
||
|
|
raise ValueError("module type not supported:", type(module1), " ", type(module2))
|
||
|
|
|
||
|
|
weight1 = get_module_weight(module1)
|
||
|
|
weight2 = get_module_weight(module2)
|
||
|
|
|
||
|
|
if weight1.size(output_axis) != weight2.size(input_axis):
|
||
|
|
raise TypeError("Number of output channels of first arg do not match \
|
||
|
|
number input channels of second arg")
|
||
|
|
|
||
|
|
bias = get_module_bias(module1)
|
||
|
|
|
||
|
|
weight1_range = channel_range(weight1, output_axis)
|
||
|
|
weight2_range = channel_range(weight2, input_axis)
|
||
|
|
|
||
|
|
# producing scaling factors to applied
|
||
|
|
weight2_range += 1e-9
|
||
|
|
scaling_factors = torch.sqrt(weight1_range / weight2_range)
|
||
|
|
inverse_scaling_factors = torch.reciprocal(scaling_factors)
|
||
|
|
|
||
|
|
bias = bias * inverse_scaling_factors
|
||
|
|
|
||
|
|
# formatting the scaling (1D) tensors to be applied on the given argument tensors
|
||
|
|
# pads axis to (1D) tensors to then be broadcasted
|
||
|
|
size1 = [1] * weight1.ndim
|
||
|
|
size1[output_axis] = weight1.size(output_axis)
|
||
|
|
size2 = [1] * weight2.ndim
|
||
|
|
size2[input_axis] = weight2.size(input_axis)
|
||
|
|
|
||
|
|
scaling_factors = torch.reshape(scaling_factors, size2)
|
||
|
|
inverse_scaling_factors = torch.reshape(inverse_scaling_factors, size1)
|
||
|
|
|
||
|
|
weight1 = weight1 * inverse_scaling_factors
|
||
|
|
weight2 = weight2 * scaling_factors
|
||
|
|
|
||
|
|
set_module_weight(module1, weight1)
|
||
|
|
set_module_bias(module1, bias)
|
||
|
|
set_module_weight(module2, weight2)
|
||
|
|
|
||
|
|
def equalize(model, paired_modules_list, threshold=1e-4, inplace=True):
|
||
|
|
"""Equalize modules until convergence is achieved.
|
||
|
|
|
||
|
|
Given a list of adjacent modules within a model, equalization will
|
||
|
|
be applied between each pair, this will repeated until convergence is achieved
|
||
|
|
|
||
|
|
Keeps a copy of the changing modules from the previous iteration, if the copies
|
||
|
|
are not that different than the current modules (determined by converged_test),
|
||
|
|
then the modules have converged enough that further equalizing is not necessary
|
||
|
|
|
||
|
|
Implementation of this referced section 4.1 of this paper https://arxiv.org/pdf/1906.04721.pdf
|
||
|
|
|
||
|
|
Args:
|
||
|
|
model: a model (nn.module) that equalization is to be applied on
|
||
|
|
paired_modules_list: a list of lists where each sublist is a pair of two
|
||
|
|
submodules found in the model, for each pair the two submodules generally
|
||
|
|
have to be adjacent in the model to get expected/reasonable results
|
||
|
|
threshold: a number used by the converged function to determine what degree
|
||
|
|
similarity between models is necessary for them to be called equivalent
|
||
|
|
inplace: determines if function is inplace or not
|
||
|
|
"""
|
||
|
|
if not inplace:
|
||
|
|
model = copy.deepcopy(model)
|
||
|
|
|
||
|
|
name_to_module : Dict[str, torch.nn.Module] = {}
|
||
|
|
previous_name_to_module: Dict[str, Any] = {}
|
||
|
|
name_set = {name for pair in paired_modules_list for name in pair}
|
||
|
|
|
||
|
|
for name, module in model.named_modules():
|
||
|
|
if name in name_set:
|
||
|
|
name_to_module[name] = module
|
||
|
|
previous_name_to_module[name] = None
|
||
|
|
while not converged(name_to_module, previous_name_to_module, threshold):
|
||
|
|
for pair in paired_modules_list:
|
||
|
|
previous_name_to_module[pair[0]] = copy.deepcopy(name_to_module[pair[0]])
|
||
|
|
previous_name_to_module[pair[1]] = copy.deepcopy(name_to_module[pair[1]])
|
||
|
|
|
||
|
|
cross_layer_equalization(name_to_module[pair[0]], name_to_module[pair[1]])
|
||
|
|
|
||
|
|
return model
|
||
|
|
|
||
|
|
def converged(curr_modules, prev_modules, threshold=1e-4):
|
||
|
|
"""Test whether modules are converged to a specified threshold.
|
||
|
|
|
||
|
|
Tests for the summed norm of the differences between each set of modules
|
||
|
|
being less than the given threshold
|
||
|
|
|
||
|
|
Takes two dictionaries mapping names to modules, the set of names for each dictionary
|
||
|
|
should be the same, looping over the set of names, for each name take the difference
|
||
|
|
between the associated modules in each dictionary
|
||
|
|
|
||
|
|
"""
|
||
|
|
if curr_modules.keys() != prev_modules.keys():
|
||
|
|
raise ValueError("The keys to the given mappings must have the same set of names of modules")
|
||
|
|
|
||
|
|
summed_norms = torch.tensor(0.)
|
||
|
|
if None in prev_modules.values():
|
||
|
|
return False
|
||
|
|
for name in curr_modules.keys():
|
||
|
|
curr_weight = get_module_weight(curr_modules[name])
|
||
|
|
prev_weight = get_module_weight(prev_modules[name])
|
||
|
|
|
||
|
|
difference = curr_weight.sub(prev_weight)
|
||
|
|
summed_norms += torch.norm(difference)
|
||
|
|
return bool(summed_norms < threshold)
|