File size: 6,981 Bytes
d1ceb73 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 |
# mypy: allow-untyped-defs
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)
|