Update q-config and black for procedures/utils
This commit is contained in:
@@ -1,7 +1,7 @@
|
||||
from .evaluation_utils import obtain_accuracy
|
||||
from .gpu_manager import GPUManager
|
||||
from .flop_benchmark import get_model_infos, count_parameters, count_parameters_in_MB
|
||||
from .affine_utils import normalize_points, denormalize_points
|
||||
from .affine_utils import identity2affine, solve2theta, affine2image
|
||||
from .hash_utils import get_md5_file
|
||||
from .str_utils import split_str2indexes
|
||||
from .gpu_manager import GPUManager
|
||||
from .flop_benchmark import get_model_infos, count_parameters, count_parameters_in_MB
|
||||
from .affine_utils import normalize_points, denormalize_points
|
||||
from .affine_utils import identity2affine, solve2theta, affine2image
|
||||
from .hash_utils import get_md5_file
|
||||
from .str_utils import split_str2indexes
|
||||
|
@@ -1,125 +1,149 @@
|
||||
# functions for affine transformation
|
||||
import math, torch
|
||||
import math
|
||||
import torch
|
||||
import numpy as np
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def identity2affine(full=False):
|
||||
if not full:
|
||||
parameters = torch.zeros((2,3))
|
||||
parameters[0, 0] = parameters[1, 1] = 1
|
||||
else:
|
||||
parameters = torch.zeros((3,3))
|
||||
parameters[0, 0] = parameters[1, 1] = parameters[2, 2] = 1
|
||||
return parameters
|
||||
if not full:
|
||||
parameters = torch.zeros((2, 3))
|
||||
parameters[0, 0] = parameters[1, 1] = 1
|
||||
else:
|
||||
parameters = torch.zeros((3, 3))
|
||||
parameters[0, 0] = parameters[1, 1] = parameters[2, 2] = 1
|
||||
return parameters
|
||||
|
||||
|
||||
def normalize_L(x, L):
|
||||
return -1. + 2. * x / (L-1)
|
||||
return -1.0 + 2.0 * x / (L - 1)
|
||||
|
||||
|
||||
def denormalize_L(x, L):
|
||||
return (x + 1.0) / 2.0 * (L-1)
|
||||
return (x + 1.0) / 2.0 * (L - 1)
|
||||
|
||||
|
||||
def crop2affine(crop_box, W, H):
|
||||
assert len(crop_box) == 4, 'Invalid crop-box : {:}'.format(crop_box)
|
||||
parameters = torch.zeros(3,3)
|
||||
x1, y1 = normalize_L(crop_box[0], W), normalize_L(crop_box[1], H)
|
||||
x2, y2 = normalize_L(crop_box[2], W), normalize_L(crop_box[3], H)
|
||||
parameters[0,0] = (x2-x1)/2
|
||||
parameters[0,2] = (x2+x1)/2
|
||||
assert len(crop_box) == 4, "Invalid crop-box : {:}".format(crop_box)
|
||||
parameters = torch.zeros(3, 3)
|
||||
x1, y1 = normalize_L(crop_box[0], W), normalize_L(crop_box[1], H)
|
||||
x2, y2 = normalize_L(crop_box[2], W), normalize_L(crop_box[3], H)
|
||||
parameters[0, 0] = (x2 - x1) / 2
|
||||
parameters[0, 2] = (x2 + x1) / 2
|
||||
|
||||
parameters[1, 1] = (y2 - y1) / 2
|
||||
parameters[1, 2] = (y2 + y1) / 2
|
||||
parameters[2, 2] = 1
|
||||
return parameters
|
||||
|
||||
parameters[1,1] = (y2-y1)/2
|
||||
parameters[1,2] = (y2+y1)/2
|
||||
parameters[2,2] = 1
|
||||
return parameters
|
||||
|
||||
def scale2affine(scalex, scaley):
|
||||
parameters = torch.zeros(3,3)
|
||||
parameters[0,0] = scalex
|
||||
parameters[1,1] = scaley
|
||||
parameters[2,2] = 1
|
||||
return parameters
|
||||
|
||||
parameters = torch.zeros(3, 3)
|
||||
parameters[0, 0] = scalex
|
||||
parameters[1, 1] = scaley
|
||||
parameters[2, 2] = 1
|
||||
return parameters
|
||||
|
||||
|
||||
def offset2affine(offx, offy):
|
||||
parameters = torch.zeros(3,3)
|
||||
parameters[0,0] = parameters[1,1] = parameters[2,2] = 1
|
||||
parameters[0,2] = offx
|
||||
parameters[1,2] = offy
|
||||
return parameters
|
||||
parameters = torch.zeros(3, 3)
|
||||
parameters[0, 0] = parameters[1, 1] = parameters[2, 2] = 1
|
||||
parameters[0, 2] = offx
|
||||
parameters[1, 2] = offy
|
||||
return parameters
|
||||
|
||||
|
||||
def horizontalmirror2affine():
|
||||
parameters = torch.zeros(3,3)
|
||||
parameters[0,0] = -1
|
||||
parameters[1,1] = parameters[2,2] = 1
|
||||
return parameters
|
||||
parameters = torch.zeros(3, 3)
|
||||
parameters[0, 0] = -1
|
||||
parameters[1, 1] = parameters[2, 2] = 1
|
||||
return parameters
|
||||
|
||||
|
||||
# clockwise rotate image = counterclockwise rotate the rectangle
|
||||
# degree is between [0, 360]
|
||||
def rotate2affine(degree):
|
||||
assert degree >= 0 and degree <= 360, 'Invalid degree : {:}'.format(degree)
|
||||
degree = degree / 180 * math.pi
|
||||
parameters = torch.zeros(3,3)
|
||||
parameters[0,0] = math.cos(-degree)
|
||||
parameters[0,1] = -math.sin(-degree)
|
||||
parameters[1,0] = math.sin(-degree)
|
||||
parameters[1,1] = math.cos(-degree)
|
||||
parameters[2,2] = 1
|
||||
return parameters
|
||||
assert degree >= 0 and degree <= 360, "Invalid degree : {:}".format(degree)
|
||||
degree = degree / 180 * math.pi
|
||||
parameters = torch.zeros(3, 3)
|
||||
parameters[0, 0] = math.cos(-degree)
|
||||
parameters[0, 1] = -math.sin(-degree)
|
||||
parameters[1, 0] = math.sin(-degree)
|
||||
parameters[1, 1] = math.cos(-degree)
|
||||
parameters[2, 2] = 1
|
||||
return parameters
|
||||
|
||||
|
||||
# shape is a tuple [H, W]
|
||||
def normalize_points(shape, points):
|
||||
assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(shape) == 2, 'invalid shape : {:}'.format(shape)
|
||||
assert isinstance(points, torch.Tensor) and (points.shape[0] == 2), 'points are wrong : {:}'.format(points.shape)
|
||||
(H, W), points = shape, points.clone()
|
||||
points[0, :] = normalize_L(points[0,:], W)
|
||||
points[1, :] = normalize_L(points[1,:], H)
|
||||
return points
|
||||
assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(shape) == 2, "invalid shape : {:}".format(
|
||||
shape
|
||||
)
|
||||
assert isinstance(points, torch.Tensor) and (points.shape[0] == 2), "points are wrong : {:}".format(points.shape)
|
||||
(H, W), points = shape, points.clone()
|
||||
points[0, :] = normalize_L(points[0, :], W)
|
||||
points[1, :] = normalize_L(points[1, :], H)
|
||||
return points
|
||||
|
||||
|
||||
# shape is a tuple [H, W]
|
||||
def normalize_points_batch(shape, points):
|
||||
assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(shape) == 2, 'invalid shape : {:}'.format(shape)
|
||||
assert isinstance(points, torch.Tensor) and (points.size(-1) == 2), 'points are wrong : {:}'.format(points.shape)
|
||||
(H, W), points = shape, points.clone()
|
||||
x = normalize_L(points[...,0], W)
|
||||
y = normalize_L(points[...,1], H)
|
||||
return torch.stack((x,y), dim=-1)
|
||||
assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(shape) == 2, "invalid shape : {:}".format(
|
||||
shape
|
||||
)
|
||||
assert isinstance(points, torch.Tensor) and (points.size(-1) == 2), "points are wrong : {:}".format(points.shape)
|
||||
(H, W), points = shape, points.clone()
|
||||
x = normalize_L(points[..., 0], W)
|
||||
y = normalize_L(points[..., 1], H)
|
||||
return torch.stack((x, y), dim=-1)
|
||||
|
||||
|
||||
# shape is a tuple [H, W]
|
||||
def denormalize_points(shape, points):
|
||||
assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(shape) == 2, 'invalid shape : {:}'.format(shape)
|
||||
assert isinstance(points, torch.Tensor) and (points.shape[0] == 2), 'points are wrong : {:}'.format(points.shape)
|
||||
(H, W), points = shape, points.clone()
|
||||
points[0, :] = denormalize_L(points[0,:], W)
|
||||
points[1, :] = denormalize_L(points[1,:], H)
|
||||
return points
|
||||
assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(shape) == 2, "invalid shape : {:}".format(
|
||||
shape
|
||||
)
|
||||
assert isinstance(points, torch.Tensor) and (points.shape[0] == 2), "points are wrong : {:}".format(points.shape)
|
||||
(H, W), points = shape, points.clone()
|
||||
points[0, :] = denormalize_L(points[0, :], W)
|
||||
points[1, :] = denormalize_L(points[1, :], H)
|
||||
return points
|
||||
|
||||
|
||||
# shape is a tuple [H, W]
|
||||
def denormalize_points_batch(shape, points):
|
||||
assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(shape) == 2, 'invalid shape : {:}'.format(shape)
|
||||
assert isinstance(points, torch.Tensor) and (points.shape[-1] == 2), 'points are wrong : {:}'.format(points.shape)
|
||||
(H, W), points = shape, points.clone()
|
||||
x = denormalize_L(points[...,0], W)
|
||||
y = denormalize_L(points[...,1], H)
|
||||
return torch.stack((x,y), dim=-1)
|
||||
assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(shape) == 2, "invalid shape : {:}".format(
|
||||
shape
|
||||
)
|
||||
assert isinstance(points, torch.Tensor) and (points.shape[-1] == 2), "points are wrong : {:}".format(points.shape)
|
||||
(H, W), points = shape, points.clone()
|
||||
x = denormalize_L(points[..., 0], W)
|
||||
y = denormalize_L(points[..., 1], H)
|
||||
return torch.stack((x, y), dim=-1)
|
||||
|
||||
|
||||
# make target * theta = source
|
||||
def solve2theta(source, target):
|
||||
source, target = source.clone(), target.clone()
|
||||
oks = source[2, :] == 1
|
||||
assert torch.sum(oks).item() >= 3, 'valid points : {:} is short'.format(oks)
|
||||
if target.size(0) == 2: target = torch.cat((target, oks.unsqueeze(0).float()), dim=0)
|
||||
source, target = source[:, oks], target[:, oks]
|
||||
source, target = source.transpose(1,0), target.transpose(1,0)
|
||||
assert source.size(1) == target.size(1) == 3
|
||||
#X, residual, rank, s = np.linalg.lstsq(target.numpy(), source.numpy())
|
||||
#theta = torch.Tensor(X.T[:2, :])
|
||||
X_, qr = torch.gels(source, target)
|
||||
theta = X_[:3, :2].transpose(1, 0)
|
||||
return theta
|
||||
source, target = source.clone(), target.clone()
|
||||
oks = source[2, :] == 1
|
||||
assert torch.sum(oks).item() >= 3, "valid points : {:} is short".format(oks)
|
||||
if target.size(0) == 2:
|
||||
target = torch.cat((target, oks.unsqueeze(0).float()), dim=0)
|
||||
source, target = source[:, oks], target[:, oks]
|
||||
source, target = source.transpose(1, 0), target.transpose(1, 0)
|
||||
assert source.size(1) == target.size(1) == 3
|
||||
# X, residual, rank, s = np.linalg.lstsq(target.numpy(), source.numpy())
|
||||
# theta = torch.Tensor(X.T[:2, :])
|
||||
X_, qr = torch.gels(source, target)
|
||||
theta = X_[:3, :2].transpose(1, 0)
|
||||
return theta
|
||||
|
||||
|
||||
# shape = [H,W]
|
||||
def affine2image(image, theta, shape):
|
||||
C, H, W = image.size()
|
||||
theta = theta[:2, :].unsqueeze(0)
|
||||
grid_size = torch.Size([1, C, shape[0], shape[1]])
|
||||
grid = F.affine_grid(theta, grid_size)
|
||||
affI = F.grid_sample(image.unsqueeze(0), grid, mode='bilinear', padding_mode='border')
|
||||
return affI.squeeze(0)
|
||||
C, H, W = image.size()
|
||||
theta = theta[:2, :].unsqueeze(0)
|
||||
grid_size = torch.Size([1, C, shape[0], shape[1]])
|
||||
grid = F.affine_grid(theta, grid_size)
|
||||
affI = F.grid_sample(image.unsqueeze(0), grid, mode="bilinear", padding_mode="border")
|
||||
return affI.squeeze(0)
|
||||
|
@@ -1,16 +1,17 @@
|
||||
import torch
|
||||
|
||||
|
||||
def obtain_accuracy(output, target, topk=(1,)):
|
||||
"""Computes the precision@k for the specified values of k"""
|
||||
maxk = max(topk)
|
||||
batch_size = target.size(0)
|
||||
"""Computes the precision@k for the specified values of k"""
|
||||
maxk = max(topk)
|
||||
batch_size = target.size(0)
|
||||
|
||||
_, pred = output.topk(maxk, 1, True, True)
|
||||
pred = pred.t()
|
||||
correct = pred.eq(target.view(1, -1).expand_as(pred))
|
||||
_, pred = output.topk(maxk, 1, True, True)
|
||||
pred = pred.t()
|
||||
correct = pred.eq(target.view(1, -1).expand_as(pred))
|
||||
|
||||
res = []
|
||||
for k in topk:
|
||||
correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
|
||||
res.append(correct_k.mul_(100.0 / batch_size))
|
||||
return res
|
||||
res = []
|
||||
for k in topk:
|
||||
correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
|
||||
res.append(correct_k.mul_(100.0 / batch_size))
|
||||
return res
|
||||
|
@@ -4,191 +4,199 @@ import numpy as np
|
||||
|
||||
|
||||
def count_parameters_in_MB(model):
|
||||
return count_parameters(model, "mb")
|
||||
return count_parameters(model, "mb")
|
||||
|
||||
|
||||
def count_parameters(model_or_parameters, unit="mb"):
|
||||
if isinstance(model_or_parameters, nn.Module):
|
||||
counts = np.sum(np.prod(v.size()) for v in model_or_parameters.parameters())
|
||||
else:
|
||||
counts = np.sum(np.prod(v.size()) for v in model_or_parameters)
|
||||
if unit.lower() == "mb":
|
||||
counts /= 1e6
|
||||
elif unit.lower() == "kb":
|
||||
counts /= 1e3
|
||||
elif unit.lower() == "gb":
|
||||
counts /= 1e9
|
||||
elif unit is not None:
|
||||
raise ValueError("Unknow unit: {:}".format(unit))
|
||||
return counts
|
||||
if isinstance(model_or_parameters, nn.Module):
|
||||
counts = np.sum(np.prod(v.size()) for v in model_or_parameters.parameters())
|
||||
else:
|
||||
counts = np.sum(np.prod(v.size()) for v in model_or_parameters)
|
||||
if unit.lower() == "mb":
|
||||
counts /= 1e6
|
||||
elif unit.lower() == "kb":
|
||||
counts /= 1e3
|
||||
elif unit.lower() == "gb":
|
||||
counts /= 1e9
|
||||
elif unit is not None:
|
||||
raise ValueError("Unknow unit: {:}".format(unit))
|
||||
return counts
|
||||
|
||||
|
||||
def get_model_infos(model, shape):
|
||||
#model = copy.deepcopy( model )
|
||||
# model = copy.deepcopy( model )
|
||||
|
||||
model = add_flops_counting_methods(model)
|
||||
#model = model.cuda()
|
||||
model.eval()
|
||||
model = add_flops_counting_methods(model)
|
||||
# model = model.cuda()
|
||||
model.eval()
|
||||
|
||||
#cache_inputs = torch.zeros(*shape).cuda()
|
||||
#cache_inputs = torch.zeros(*shape)
|
||||
cache_inputs = torch.rand(*shape)
|
||||
if next(model.parameters()).is_cuda: cache_inputs = cache_inputs.cuda()
|
||||
#print_log('In the calculating function : cache input size : {:}'.format(cache_inputs.size()), log)
|
||||
with torch.no_grad():
|
||||
_____ = model(cache_inputs)
|
||||
FLOPs = compute_average_flops_cost( model ) / 1e6
|
||||
Param = count_parameters_in_MB(model)
|
||||
# cache_inputs = torch.zeros(*shape).cuda()
|
||||
# cache_inputs = torch.zeros(*shape)
|
||||
cache_inputs = torch.rand(*shape)
|
||||
if next(model.parameters()).is_cuda:
|
||||
cache_inputs = cache_inputs.cuda()
|
||||
# print_log('In the calculating function : cache input size : {:}'.format(cache_inputs.size()), log)
|
||||
with torch.no_grad():
|
||||
_____ = model(cache_inputs)
|
||||
FLOPs = compute_average_flops_cost(model) / 1e6
|
||||
Param = count_parameters_in_MB(model)
|
||||
|
||||
if hasattr(model, 'auxiliary_param'):
|
||||
aux_params = count_parameters_in_MB(model.auxiliary_param())
|
||||
print ('The auxiliary params of this model is : {:}'.format(aux_params))
|
||||
print ('We remove the auxiliary params from the total params ({:}) when counting'.format(Param))
|
||||
Param = Param - aux_params
|
||||
|
||||
#print_log('FLOPs : {:} MB'.format(FLOPs), log)
|
||||
torch.cuda.empty_cache()
|
||||
model.apply( remove_hook_function )
|
||||
return FLOPs, Param
|
||||
if hasattr(model, "auxiliary_param"):
|
||||
aux_params = count_parameters_in_MB(model.auxiliary_param())
|
||||
print("The auxiliary params of this model is : {:}".format(aux_params))
|
||||
print("We remove the auxiliary params from the total params ({:}) when counting".format(Param))
|
||||
Param = Param - aux_params
|
||||
|
||||
# print_log('FLOPs : {:} MB'.format(FLOPs), log)
|
||||
torch.cuda.empty_cache()
|
||||
model.apply(remove_hook_function)
|
||||
return FLOPs, Param
|
||||
|
||||
|
||||
# ---- Public functions
|
||||
def add_flops_counting_methods( model ):
|
||||
model.__batch_counter__ = 0
|
||||
add_batch_counter_hook_function( model )
|
||||
model.apply( add_flops_counter_variable_or_reset )
|
||||
model.apply( add_flops_counter_hook_function )
|
||||
return model
|
||||
|
||||
def add_flops_counting_methods(model):
|
||||
model.__batch_counter__ = 0
|
||||
add_batch_counter_hook_function(model)
|
||||
model.apply(add_flops_counter_variable_or_reset)
|
||||
model.apply(add_flops_counter_hook_function)
|
||||
return model
|
||||
|
||||
|
||||
def compute_average_flops_cost(model):
|
||||
"""
|
||||
A method that will be available after add_flops_counting_methods() is called on a desired net object.
|
||||
Returns current mean flops consumption per image.
|
||||
"""
|
||||
batches_count = model.__batch_counter__
|
||||
flops_sum = 0
|
||||
#or isinstance(module, torch.nn.AvgPool2d) or isinstance(module, torch.nn.MaxPool2d) \
|
||||
for module in model.modules():
|
||||
if isinstance(module, torch.nn.Conv2d) or isinstance(module, torch.nn.Linear) \
|
||||
or isinstance(module, torch.nn.Conv1d) \
|
||||
or hasattr(module, 'calculate_flop_self'):
|
||||
flops_sum += module.__flops__
|
||||
return flops_sum / batches_count
|
||||
"""
|
||||
A method that will be available after add_flops_counting_methods() is called on a desired net object.
|
||||
Returns current mean flops consumption per image.
|
||||
"""
|
||||
batches_count = model.__batch_counter__
|
||||
flops_sum = 0
|
||||
# or isinstance(module, torch.nn.AvgPool2d) or isinstance(module, torch.nn.MaxPool2d) \
|
||||
for module in model.modules():
|
||||
if (
|
||||
isinstance(module, torch.nn.Conv2d)
|
||||
or isinstance(module, torch.nn.Linear)
|
||||
or isinstance(module, torch.nn.Conv1d)
|
||||
or hasattr(module, "calculate_flop_self")
|
||||
):
|
||||
flops_sum += module.__flops__
|
||||
return flops_sum / batches_count
|
||||
|
||||
|
||||
# ---- Internal functions
|
||||
def pool_flops_counter_hook(pool_module, inputs, output):
|
||||
batch_size = inputs[0].size(0)
|
||||
kernel_size = pool_module.kernel_size
|
||||
out_C, output_height, output_width = output.shape[1:]
|
||||
assert out_C == inputs[0].size(1), '{:} vs. {:}'.format(out_C, inputs[0].size())
|
||||
batch_size = inputs[0].size(0)
|
||||
kernel_size = pool_module.kernel_size
|
||||
out_C, output_height, output_width = output.shape[1:]
|
||||
assert out_C == inputs[0].size(1), "{:} vs. {:}".format(out_C, inputs[0].size())
|
||||
|
||||
overall_flops = batch_size * out_C * output_height * output_width * kernel_size * kernel_size
|
||||
pool_module.__flops__ += overall_flops
|
||||
overall_flops = batch_size * out_C * output_height * output_width * kernel_size * kernel_size
|
||||
pool_module.__flops__ += overall_flops
|
||||
|
||||
|
||||
def self_calculate_flops_counter_hook(self_module, inputs, output):
|
||||
overall_flops = self_module.calculate_flop_self(inputs[0].shape, output.shape)
|
||||
self_module.__flops__ += overall_flops
|
||||
overall_flops = self_module.calculate_flop_self(inputs[0].shape, output.shape)
|
||||
self_module.__flops__ += overall_flops
|
||||
|
||||
|
||||
def fc_flops_counter_hook(fc_module, inputs, output):
|
||||
batch_size = inputs[0].size(0)
|
||||
xin, xout = fc_module.in_features, fc_module.out_features
|
||||
assert xin == inputs[0].size(1) and xout == output.size(1), 'IO=({:}, {:})'.format(xin, xout)
|
||||
overall_flops = batch_size * xin * xout
|
||||
if fc_module.bias is not None:
|
||||
overall_flops += batch_size * xout
|
||||
fc_module.__flops__ += overall_flops
|
||||
batch_size = inputs[0].size(0)
|
||||
xin, xout = fc_module.in_features, fc_module.out_features
|
||||
assert xin == inputs[0].size(1) and xout == output.size(1), "IO=({:}, {:})".format(xin, xout)
|
||||
overall_flops = batch_size * xin * xout
|
||||
if fc_module.bias is not None:
|
||||
overall_flops += batch_size * xout
|
||||
fc_module.__flops__ += overall_flops
|
||||
|
||||
|
||||
def conv1d_flops_counter_hook(conv_module, inputs, outputs):
|
||||
batch_size = inputs[0].size(0)
|
||||
outL = outputs.shape[-1]
|
||||
[kernel] = conv_module.kernel_size
|
||||
in_channels = conv_module.in_channels
|
||||
out_channels = conv_module.out_channels
|
||||
groups = conv_module.groups
|
||||
conv_per_position_flops = kernel * in_channels * out_channels / groups
|
||||
|
||||
active_elements_count = batch_size * outL
|
||||
overall_flops = conv_per_position_flops * active_elements_count
|
||||
batch_size = inputs[0].size(0)
|
||||
outL = outputs.shape[-1]
|
||||
[kernel] = conv_module.kernel_size
|
||||
in_channels = conv_module.in_channels
|
||||
out_channels = conv_module.out_channels
|
||||
groups = conv_module.groups
|
||||
conv_per_position_flops = kernel * in_channels * out_channels / groups
|
||||
|
||||
if conv_module.bias is not None:
|
||||
overall_flops += out_channels * active_elements_count
|
||||
conv_module.__flops__ += overall_flops
|
||||
active_elements_count = batch_size * outL
|
||||
overall_flops = conv_per_position_flops * active_elements_count
|
||||
|
||||
if conv_module.bias is not None:
|
||||
overall_flops += out_channels * active_elements_count
|
||||
conv_module.__flops__ += overall_flops
|
||||
|
||||
|
||||
def conv2d_flops_counter_hook(conv_module, inputs, output):
|
||||
batch_size = inputs[0].size(0)
|
||||
output_height, output_width = output.shape[2:]
|
||||
|
||||
kernel_height, kernel_width = conv_module.kernel_size
|
||||
in_channels = conv_module.in_channels
|
||||
out_channels = conv_module.out_channels
|
||||
groups = conv_module.groups
|
||||
conv_per_position_flops = kernel_height * kernel_width * in_channels * out_channels / groups
|
||||
|
||||
active_elements_count = batch_size * output_height * output_width
|
||||
overall_flops = conv_per_position_flops * active_elements_count
|
||||
|
||||
if conv_module.bias is not None:
|
||||
overall_flops += out_channels * active_elements_count
|
||||
conv_module.__flops__ += overall_flops
|
||||
batch_size = inputs[0].size(0)
|
||||
output_height, output_width = output.shape[2:]
|
||||
|
||||
kernel_height, kernel_width = conv_module.kernel_size
|
||||
in_channels = conv_module.in_channels
|
||||
out_channels = conv_module.out_channels
|
||||
groups = conv_module.groups
|
||||
conv_per_position_flops = kernel_height * kernel_width * in_channels * out_channels / groups
|
||||
|
||||
active_elements_count = batch_size * output_height * output_width
|
||||
overall_flops = conv_per_position_flops * active_elements_count
|
||||
|
||||
if conv_module.bias is not None:
|
||||
overall_flops += out_channels * active_elements_count
|
||||
conv_module.__flops__ += overall_flops
|
||||
|
||||
|
||||
|
||||
def batch_counter_hook(module, inputs, output):
|
||||
# Can have multiple inputs, getting the first one
|
||||
inputs = inputs[0]
|
||||
batch_size = inputs.shape[0]
|
||||
module.__batch_counter__ += batch_size
|
||||
# Can have multiple inputs, getting the first one
|
||||
inputs = inputs[0]
|
||||
batch_size = inputs.shape[0]
|
||||
module.__batch_counter__ += batch_size
|
||||
|
||||
|
||||
def add_batch_counter_hook_function(module):
|
||||
if not hasattr(module, '__batch_counter_handle__'):
|
||||
handle = module.register_forward_hook(batch_counter_hook)
|
||||
module.__batch_counter_handle__ = handle
|
||||
if not hasattr(module, "__batch_counter_handle__"):
|
||||
handle = module.register_forward_hook(batch_counter_hook)
|
||||
module.__batch_counter_handle__ = handle
|
||||
|
||||
|
||||
|
||||
def add_flops_counter_variable_or_reset(module):
|
||||
if isinstance(module, torch.nn.Conv2d) or isinstance(module, torch.nn.Linear) \
|
||||
or isinstance(module, torch.nn.Conv1d) \
|
||||
or isinstance(module, torch.nn.AvgPool2d) or isinstance(module, torch.nn.MaxPool2d) \
|
||||
or hasattr(module, 'calculate_flop_self'):
|
||||
module.__flops__ = 0
|
||||
if (
|
||||
isinstance(module, torch.nn.Conv2d)
|
||||
or isinstance(module, torch.nn.Linear)
|
||||
or isinstance(module, torch.nn.Conv1d)
|
||||
or isinstance(module, torch.nn.AvgPool2d)
|
||||
or isinstance(module, torch.nn.MaxPool2d)
|
||||
or hasattr(module, "calculate_flop_self")
|
||||
):
|
||||
module.__flops__ = 0
|
||||
|
||||
|
||||
def add_flops_counter_hook_function(module):
|
||||
if isinstance(module, torch.nn.Conv2d):
|
||||
if not hasattr(module, '__flops_handle__'):
|
||||
handle = module.register_forward_hook(conv2d_flops_counter_hook)
|
||||
module.__flops_handle__ = handle
|
||||
elif isinstance(module, torch.nn.Conv1d):
|
||||
if not hasattr(module, '__flops_handle__'):
|
||||
handle = module.register_forward_hook(conv1d_flops_counter_hook)
|
||||
module.__flops_handle__ = handle
|
||||
elif isinstance(module, torch.nn.Linear):
|
||||
if not hasattr(module, '__flops_handle__'):
|
||||
handle = module.register_forward_hook(fc_flops_counter_hook)
|
||||
module.__flops_handle__ = handle
|
||||
elif isinstance(module, torch.nn.AvgPool2d) or isinstance(module, torch.nn.MaxPool2d):
|
||||
if not hasattr(module, '__flops_handle__'):
|
||||
handle = module.register_forward_hook(pool_flops_counter_hook)
|
||||
module.__flops_handle__ = handle
|
||||
elif hasattr(module, 'calculate_flop_self'): # self-defined module
|
||||
if not hasattr(module, '__flops_handle__'):
|
||||
handle = module.register_forward_hook(self_calculate_flops_counter_hook)
|
||||
module.__flops_handle__ = handle
|
||||
if isinstance(module, torch.nn.Conv2d):
|
||||
if not hasattr(module, "__flops_handle__"):
|
||||
handle = module.register_forward_hook(conv2d_flops_counter_hook)
|
||||
module.__flops_handle__ = handle
|
||||
elif isinstance(module, torch.nn.Conv1d):
|
||||
if not hasattr(module, "__flops_handle__"):
|
||||
handle = module.register_forward_hook(conv1d_flops_counter_hook)
|
||||
module.__flops_handle__ = handle
|
||||
elif isinstance(module, torch.nn.Linear):
|
||||
if not hasattr(module, "__flops_handle__"):
|
||||
handle = module.register_forward_hook(fc_flops_counter_hook)
|
||||
module.__flops_handle__ = handle
|
||||
elif isinstance(module, torch.nn.AvgPool2d) or isinstance(module, torch.nn.MaxPool2d):
|
||||
if not hasattr(module, "__flops_handle__"):
|
||||
handle = module.register_forward_hook(pool_flops_counter_hook)
|
||||
module.__flops_handle__ = handle
|
||||
elif hasattr(module, "calculate_flop_self"): # self-defined module
|
||||
if not hasattr(module, "__flops_handle__"):
|
||||
handle = module.register_forward_hook(self_calculate_flops_counter_hook)
|
||||
module.__flops_handle__ = handle
|
||||
|
||||
|
||||
def remove_hook_function(module):
|
||||
hookers = ['__batch_counter_handle__', '__flops_handle__']
|
||||
for hooker in hookers:
|
||||
if hasattr(module, hooker):
|
||||
handle = getattr(module, hooker)
|
||||
handle.remove()
|
||||
keys = ['__flops__', '__batch_counter__', '__flops__'] + hookers
|
||||
for ckey in keys:
|
||||
if hasattr(module, ckey): delattr(module, ckey)
|
||||
hookers = ["__batch_counter_handle__", "__flops_handle__"]
|
||||
for hooker in hookers:
|
||||
if hasattr(module, hooker):
|
||||
handle = getattr(module, hooker)
|
||||
handle.remove()
|
||||
keys = ["__flops__", "__batch_counter__", "__flops__"] + hookers
|
||||
for ckey in keys:
|
||||
if hasattr(module, ckey):
|
||||
delattr(module, ckey)
|
||||
|
@@ -1,65 +1,69 @@
|
||||
import os
|
||||
|
||||
class GPUManager():
|
||||
queries = ('index', 'gpu_name', 'memory.free', 'memory.used', 'memory.total', 'power.draw', 'power.limit')
|
||||
|
||||
def __init__(self):
|
||||
all_gpus = self.query_gpu(False)
|
||||
class GPUManager:
|
||||
queries = ("index", "gpu_name", "memory.free", "memory.used", "memory.total", "power.draw", "power.limit")
|
||||
|
||||
def get_info(self, ctype):
|
||||
cmd = 'nvidia-smi --query-gpu={} --format=csv,noheader'.format(ctype)
|
||||
lines = os.popen(cmd).readlines()
|
||||
lines = [line.strip('\n') for line in lines]
|
||||
return lines
|
||||
def __init__(self):
|
||||
all_gpus = self.query_gpu(False)
|
||||
|
||||
def query_gpu(self, show=True):
|
||||
num_gpus = len( self.get_info('index') )
|
||||
all_gpus = [ {} for i in range(num_gpus) ]
|
||||
for query in self.queries:
|
||||
infos = self.get_info(query)
|
||||
for idx, info in enumerate(infos):
|
||||
all_gpus[idx][query] = info
|
||||
def get_info(self, ctype):
|
||||
cmd = "nvidia-smi --query-gpu={} --format=csv,noheader".format(ctype)
|
||||
lines = os.popen(cmd).readlines()
|
||||
lines = [line.strip("\n") for line in lines]
|
||||
return lines
|
||||
|
||||
if 'CUDA_VISIBLE_DEVICES' in os.environ:
|
||||
CUDA_VISIBLE_DEVICES = os.environ['CUDA_VISIBLE_DEVICES'].split(',')
|
||||
selected_gpus = []
|
||||
for idx, CUDA_VISIBLE_DEVICE in enumerate(CUDA_VISIBLE_DEVICES):
|
||||
find = False
|
||||
for gpu in all_gpus:
|
||||
if gpu['index'] == CUDA_VISIBLE_DEVICE:
|
||||
assert not find, 'Duplicate cuda device index : {}'.format(CUDA_VISIBLE_DEVICE)
|
||||
find = True
|
||||
selected_gpus.append( gpu.copy() )
|
||||
selected_gpus[-1]['index'] = '{}'.format(idx)
|
||||
assert find, 'Does not find the device : {}'.format(CUDA_VISIBLE_DEVICE)
|
||||
all_gpus = selected_gpus
|
||||
|
||||
if show:
|
||||
allstrings = ''
|
||||
for gpu in all_gpus:
|
||||
string = '| '
|
||||
def query_gpu(self, show=True):
|
||||
num_gpus = len(self.get_info("index"))
|
||||
all_gpus = [{} for i in range(num_gpus)]
|
||||
for query in self.queries:
|
||||
if query.find('memory') == 0: xinfo = '{:>9}'.format(gpu[query])
|
||||
else: xinfo = gpu[query]
|
||||
string = string + query + ' : ' + xinfo + ' | '
|
||||
allstrings = allstrings + string + '\n'
|
||||
return allstrings
|
||||
else:
|
||||
return all_gpus
|
||||
infos = self.get_info(query)
|
||||
for idx, info in enumerate(infos):
|
||||
all_gpus[idx][query] = info
|
||||
|
||||
if "CUDA_VISIBLE_DEVICES" in os.environ:
|
||||
CUDA_VISIBLE_DEVICES = os.environ["CUDA_VISIBLE_DEVICES"].split(",")
|
||||
selected_gpus = []
|
||||
for idx, CUDA_VISIBLE_DEVICE in enumerate(CUDA_VISIBLE_DEVICES):
|
||||
find = False
|
||||
for gpu in all_gpus:
|
||||
if gpu["index"] == CUDA_VISIBLE_DEVICE:
|
||||
assert not find, "Duplicate cuda device index : {}".format(CUDA_VISIBLE_DEVICE)
|
||||
find = True
|
||||
selected_gpus.append(gpu.copy())
|
||||
selected_gpus[-1]["index"] = "{}".format(idx)
|
||||
assert find, "Does not find the device : {}".format(CUDA_VISIBLE_DEVICE)
|
||||
all_gpus = selected_gpus
|
||||
|
||||
if show:
|
||||
allstrings = ""
|
||||
for gpu in all_gpus:
|
||||
string = "| "
|
||||
for query in self.queries:
|
||||
if query.find("memory") == 0:
|
||||
xinfo = "{:>9}".format(gpu[query])
|
||||
else:
|
||||
xinfo = gpu[query]
|
||||
string = string + query + " : " + xinfo + " | "
|
||||
allstrings = allstrings + string + "\n"
|
||||
return allstrings
|
||||
else:
|
||||
return all_gpus
|
||||
|
||||
def select_by_memory(self, numbers=1):
|
||||
all_gpus = self.query_gpu(False)
|
||||
assert numbers <= len(all_gpus), "Require {} gpus more than you have".format(numbers)
|
||||
alls = []
|
||||
for idx, gpu in enumerate(all_gpus):
|
||||
free_memory = gpu["memory.free"]
|
||||
free_memory = free_memory.split(" ")[0]
|
||||
free_memory = int(free_memory)
|
||||
index = gpu["index"]
|
||||
alls.append((free_memory, index))
|
||||
alls.sort(reverse=True)
|
||||
alls = [int(alls[i][1]) for i in range(numbers)]
|
||||
return sorted(alls)
|
||||
|
||||
def select_by_memory(self, numbers=1):
|
||||
all_gpus = self.query_gpu(False)
|
||||
assert numbers <= len(all_gpus), 'Require {} gpus more than you have'.format(numbers)
|
||||
alls = []
|
||||
for idx, gpu in enumerate(all_gpus):
|
||||
free_memory = gpu['memory.free']
|
||||
free_memory = free_memory.split(' ')[0]
|
||||
free_memory = int(free_memory)
|
||||
index = gpu['index']
|
||||
alls.append((free_memory, index))
|
||||
alls.sort(reverse = True)
|
||||
alls = [ int(alls[i][1]) for i in range(numbers) ]
|
||||
return sorted(alls)
|
||||
|
||||
"""
|
||||
if __name__ == '__main__':
|
||||
|
@@ -1,16 +1,17 @@
|
||||
import os, hashlib
|
||||
import os
|
||||
import hashlib
|
||||
|
||||
|
||||
def get_md5_file(file_path, post_truncated=5):
|
||||
md5_hash = hashlib.md5()
|
||||
if os.path.exists(file_path):
|
||||
xfile = open(file_path, "rb")
|
||||
content = xfile.read()
|
||||
md5_hash.update(content)
|
||||
digest = md5_hash.hexdigest()
|
||||
else:
|
||||
raise ValueError('[get_md5_file] {:} does not exist'.format(file_path))
|
||||
if post_truncated is None:
|
||||
return digest
|
||||
else:
|
||||
return digest[-post_truncated:]
|
||||
md5_hash = hashlib.md5()
|
||||
if os.path.exists(file_path):
|
||||
xfile = open(file_path, "rb")
|
||||
content = xfile.read()
|
||||
md5_hash.update(content)
|
||||
digest = md5_hash.hexdigest()
|
||||
else:
|
||||
raise ValueError("[get_md5_file] {:} does not exist".format(file_path))
|
||||
if post_truncated is None:
|
||||
return digest
|
||||
else:
|
||||
return digest[-post_truncated:]
|
||||
|
@@ -10,48 +10,58 @@ from log_utils import time_string
|
||||
|
||||
|
||||
def evaluate_one_shot(model, xloader, api, cal_mode, seed=111):
|
||||
print ('This is an old version of codes to use NAS-Bench-API, and should be modified to align with the new version. Please contact me for more details if you use this function.')
|
||||
weights = deepcopy(model.state_dict())
|
||||
model.train(cal_mode)
|
||||
with torch.no_grad():
|
||||
logits = nn.functional.log_softmax(model.arch_parameters, dim=-1)
|
||||
archs = CellStructure.gen_all(model.op_names, model.max_nodes, False)
|
||||
probs, accuracies, gt_accs_10_valid, gt_accs_10_test = [], [], [], []
|
||||
loader_iter = iter(xloader)
|
||||
random.seed(seed)
|
||||
random.shuffle(archs)
|
||||
for idx, arch in enumerate(archs):
|
||||
arch_index = api.query_index_by_arch( arch )
|
||||
metrics = api.get_more_info(arch_index, 'cifar10-valid', None, False, False)
|
||||
gt_accs_10_valid.append( metrics['valid-accuracy'] )
|
||||
metrics = api.get_more_info(arch_index, 'cifar10', None, False, False)
|
||||
gt_accs_10_test.append( metrics['test-accuracy'] )
|
||||
select_logits = []
|
||||
for i, node_info in enumerate(arch.nodes):
|
||||
for op, xin in node_info:
|
||||
node_str = '{:}<-{:}'.format(i+1, xin)
|
||||
op_index = model.op_names.index(op)
|
||||
select_logits.append( logits[model.edge2index[node_str], op_index] )
|
||||
cur_prob = sum(select_logits).item()
|
||||
probs.append( cur_prob )
|
||||
cor_prob_valid = np.corrcoef(probs, gt_accs_10_valid)[0,1]
|
||||
cor_prob_test = np.corrcoef(probs, gt_accs_10_test )[0,1]
|
||||
print ('{:} correlation for probabilities : {:.6f} on CIFAR-10 validation and {:.6f} on CIFAR-10 test'.format(time_string(), cor_prob_valid, cor_prob_test))
|
||||
|
||||
for idx, arch in enumerate(archs):
|
||||
model.set_cal_mode('dynamic', arch)
|
||||
try:
|
||||
inputs, targets = next(loader_iter)
|
||||
except:
|
||||
print(
|
||||
"This is an old version of codes to use NAS-Bench-API, and should be modified to align with the new version. Please contact me for more details if you use this function."
|
||||
)
|
||||
weights = deepcopy(model.state_dict())
|
||||
model.train(cal_mode)
|
||||
with torch.no_grad():
|
||||
logits = nn.functional.log_softmax(model.arch_parameters, dim=-1)
|
||||
archs = CellStructure.gen_all(model.op_names, model.max_nodes, False)
|
||||
probs, accuracies, gt_accs_10_valid, gt_accs_10_test = [], [], [], []
|
||||
loader_iter = iter(xloader)
|
||||
inputs, targets = next(loader_iter)
|
||||
_, logits = model(inputs.cuda())
|
||||
_, preds = torch.max(logits, dim=-1)
|
||||
correct = (preds == targets.cuda() ).float()
|
||||
accuracies.append( correct.mean().item() )
|
||||
if idx != 0 and (idx % 500 == 0 or idx + 1 == len(archs)):
|
||||
cor_accs_valid = np.corrcoef(accuracies, gt_accs_10_valid[:idx+1])[0,1]
|
||||
cor_accs_test = np.corrcoef(accuracies, gt_accs_10_test [:idx+1])[0,1]
|
||||
print ('{:} {:05d}/{:05d} mode={:5s}, correlation : accs={:.5f} for CIFAR-10 valid, {:.5f} for CIFAR-10 test.'.format(time_string(), idx, len(archs), 'Train' if cal_mode else 'Eval', cor_accs_valid, cor_accs_test))
|
||||
model.load_state_dict(weights)
|
||||
return archs, probs, accuracies
|
||||
random.seed(seed)
|
||||
random.shuffle(archs)
|
||||
for idx, arch in enumerate(archs):
|
||||
arch_index = api.query_index_by_arch(arch)
|
||||
metrics = api.get_more_info(arch_index, "cifar10-valid", None, False, False)
|
||||
gt_accs_10_valid.append(metrics["valid-accuracy"])
|
||||
metrics = api.get_more_info(arch_index, "cifar10", None, False, False)
|
||||
gt_accs_10_test.append(metrics["test-accuracy"])
|
||||
select_logits = []
|
||||
for i, node_info in enumerate(arch.nodes):
|
||||
for op, xin in node_info:
|
||||
node_str = "{:}<-{:}".format(i + 1, xin)
|
||||
op_index = model.op_names.index(op)
|
||||
select_logits.append(logits[model.edge2index[node_str], op_index])
|
||||
cur_prob = sum(select_logits).item()
|
||||
probs.append(cur_prob)
|
||||
cor_prob_valid = np.corrcoef(probs, gt_accs_10_valid)[0, 1]
|
||||
cor_prob_test = np.corrcoef(probs, gt_accs_10_test)[0, 1]
|
||||
print(
|
||||
"{:} correlation for probabilities : {:.6f} on CIFAR-10 validation and {:.6f} on CIFAR-10 test".format(
|
||||
time_string(), cor_prob_valid, cor_prob_test
|
||||
)
|
||||
)
|
||||
|
||||
for idx, arch in enumerate(archs):
|
||||
model.set_cal_mode("dynamic", arch)
|
||||
try:
|
||||
inputs, targets = next(loader_iter)
|
||||
except:
|
||||
loader_iter = iter(xloader)
|
||||
inputs, targets = next(loader_iter)
|
||||
_, logits = model(inputs.cuda())
|
||||
_, preds = torch.max(logits, dim=-1)
|
||||
correct = (preds == targets.cuda()).float()
|
||||
accuracies.append(correct.mean().item())
|
||||
if idx != 0 and (idx % 500 == 0 or idx + 1 == len(archs)):
|
||||
cor_accs_valid = np.corrcoef(accuracies, gt_accs_10_valid[: idx + 1])[0, 1]
|
||||
cor_accs_test = np.corrcoef(accuracies, gt_accs_10_test[: idx + 1])[0, 1]
|
||||
print(
|
||||
"{:} {:05d}/{:05d} mode={:5s}, correlation : accs={:.5f} for CIFAR-10 valid, {:.5f} for CIFAR-10 test.".format(
|
||||
time_string(), idx, len(archs), "Train" if cal_mode else "Eval", cor_accs_valid, cor_accs_test
|
||||
)
|
||||
)
|
||||
model.load_state_dict(weights)
|
||||
return archs, probs, accuracies
|
||||
|
@@ -1,18 +1,17 @@
|
||||
|
||||
def split_str2indexes(string: str, max_check: int, length_limit=5):
|
||||
if not isinstance(string, str):
|
||||
raise ValueError('Invalid scheme for {:}'.format(string))
|
||||
srangestr = "".join(string.split())
|
||||
indexes = set()
|
||||
for srange in srangestr.split(','):
|
||||
srange = srange.split('-')
|
||||
if len(srange) != 2:
|
||||
raise ValueError('invalid srange : {:}'.format(srange))
|
||||
if length_limit is not None:
|
||||
assert len(srange[0]) == len(srange[1]) == length_limit, 'invalid srange : {:}'.format(srange)
|
||||
srange = (int(srange[0]), int(srange[1]))
|
||||
if not (0 <= srange[0] <= srange[1] < max_check):
|
||||
raise ValueError('{:} vs {:} vs {:}'.format(srange[0], srange[1], max_check))
|
||||
for i in range(srange[0], srange[1]+1):
|
||||
indexes.add(i)
|
||||
return indexes
|
||||
if not isinstance(string, str):
|
||||
raise ValueError("Invalid scheme for {:}".format(string))
|
||||
srangestr = "".join(string.split())
|
||||
indexes = set()
|
||||
for srange in srangestr.split(","):
|
||||
srange = srange.split("-")
|
||||
if len(srange) != 2:
|
||||
raise ValueError("invalid srange : {:}".format(srange))
|
||||
if length_limit is not None:
|
||||
assert len(srange[0]) == len(srange[1]) == length_limit, "invalid srange : {:}".format(srange)
|
||||
srange = (int(srange[0]), int(srange[1]))
|
||||
if not (0 <= srange[0] <= srange[1] < max_check):
|
||||
raise ValueError("{:} vs {:} vs {:}".format(srange[0], srange[1], max_check))
|
||||
for i in range(srange[0], srange[1] + 1):
|
||||
indexes.add(i)
|
||||
return indexes
|
||||
|
@@ -11,309 +11,350 @@ from sklearn.decomposition import TruncatedSVD
|
||||
|
||||
|
||||
def available_module_types():
|
||||
return (nn.Conv2d, nn.Linear)
|
||||
return (nn.Conv2d, nn.Linear)
|
||||
|
||||
|
||||
def get_conv2D_Wmats(tensor: np.ndarray) -> List[np.ndarray]:
|
||||
"""
|
||||
Extract W slices from a 4 index conv2D tensor of shape: (N,M,i,j) or (M,N,i,j).
|
||||
Return ij (N x M) matrices
|
||||
"""
|
||||
mats = []
|
||||
N, M, imax, jmax = tensor.shape
|
||||
assert N + M >= imax + jmax, 'invalid tensor shape detected: {}x{} (NxM), {}x{} (i,j)'.format(N, M, imax, jmax)
|
||||
for i in range(imax):
|
||||
for j in range(jmax):
|
||||
w = tensor[:, :, i, j]
|
||||
if N < M: w = w.T
|
||||
mats.append(w)
|
||||
return mats
|
||||
"""
|
||||
Extract W slices from a 4 index conv2D tensor of shape: (N,M,i,j) or (M,N,i,j).
|
||||
Return ij (N x M) matrices
|
||||
"""
|
||||
mats = []
|
||||
N, M, imax, jmax = tensor.shape
|
||||
assert N + M >= imax + jmax, "invalid tensor shape detected: {}x{} (NxM), {}x{} (i,j)".format(N, M, imax, jmax)
|
||||
for i in range(imax):
|
||||
for j in range(jmax):
|
||||
w = tensor[:, :, i, j]
|
||||
if N < M:
|
||||
w = w.T
|
||||
mats.append(w)
|
||||
return mats
|
||||
|
||||
|
||||
def glorot_norm_check(W, N, M, rf_size, lower=0.5, upper=1.5):
|
||||
"""Check if this layer needs Glorot Normalization Fix"""
|
||||
"""Check if this layer needs Glorot Normalization Fix"""
|
||||
|
||||
kappa = np.sqrt(2 / ((N + M) * rf_size))
|
||||
norm = np.linalg.norm(W)
|
||||
kappa = np.sqrt(2 / ((N + M) * rf_size))
|
||||
norm = np.linalg.norm(W)
|
||||
|
||||
check1 = norm / np.sqrt(N * M)
|
||||
check2 = norm / (kappa * np.sqrt(N * M))
|
||||
check1 = norm / np.sqrt(N * M)
|
||||
check2 = norm / (kappa * np.sqrt(N * M))
|
||||
|
||||
if (rf_size > 1) and (check2 > lower) and (check2 < upper):
|
||||
return check2, True
|
||||
elif (check1 > lower) & (check1 < upper):
|
||||
return check1, True
|
||||
else:
|
||||
if rf_size > 1:
|
||||
return check2, False
|
||||
else:
|
||||
return check1, False
|
||||
|
||||
if (rf_size > 1) and (check2 > lower) and (check2 < upper):
|
||||
return check2, True
|
||||
elif (check1 > lower) & (check1 < upper):
|
||||
return check1, True
|
||||
else:
|
||||
if rf_size > 1: return check2, False
|
||||
else: return check1, False
|
||||
|
||||
def glorot_norm_fix(w, n, m, rf_size):
|
||||
"""Apply Glorot Normalization Fix."""
|
||||
kappa = np.sqrt(2 / ((n + m) * rf_size))
|
||||
w = w / kappa
|
||||
return w
|
||||
"""Apply Glorot Normalization Fix."""
|
||||
kappa = np.sqrt(2 / ((n + m) * rf_size))
|
||||
w = w / kappa
|
||||
return w
|
||||
|
||||
|
||||
def analyze_weights(weights, min_size, max_size, alphas, lognorms, spectralnorms, softranks, normalize, glorot_fix):
|
||||
results = OrderedDict()
|
||||
count = len(weights)
|
||||
if count == 0: return results
|
||||
results = OrderedDict()
|
||||
count = len(weights)
|
||||
if count == 0:
|
||||
return results
|
||||
|
||||
for i, weight in enumerate(weights):
|
||||
M, N = np.min(weight.shape), np.max(weight.shape)
|
||||
Q = N / M
|
||||
results[i] = cur_res = OrderedDict(N=N, M=M, Q=Q)
|
||||
check, checkTF = glorot_norm_check(weight, N, M, count)
|
||||
cur_res['check'] = check
|
||||
cur_res['checkTF'] = checkTF
|
||||
# assume receptive field size is count
|
||||
if glorot_fix:
|
||||
weight = glorot_norm_fix(weight, N, M, count)
|
||||
else:
|
||||
# probably never needed since we always fix for glorot
|
||||
weight = weight * np.sqrt(count / 2.0)
|
||||
for i, weight in enumerate(weights):
|
||||
M, N = np.min(weight.shape), np.max(weight.shape)
|
||||
Q = N / M
|
||||
results[i] = cur_res = OrderedDict(N=N, M=M, Q=Q)
|
||||
check, checkTF = glorot_norm_check(weight, N, M, count)
|
||||
cur_res["check"] = check
|
||||
cur_res["checkTF"] = checkTF
|
||||
# assume receptive field size is count
|
||||
if glorot_fix:
|
||||
weight = glorot_norm_fix(weight, N, M, count)
|
||||
else:
|
||||
# probably never needed since we always fix for glorot
|
||||
weight = weight * np.sqrt(count / 2.0)
|
||||
|
||||
if spectralnorms: # spectralnorm is the max eigenvalues
|
||||
svd = TruncatedSVD(n_components=1, n_iter=7, random_state=10)
|
||||
svd.fit(weight)
|
||||
sv = svd.singular_values_
|
||||
sv_max = np.max(sv)
|
||||
if normalize:
|
||||
evals = sv * sv / N
|
||||
else:
|
||||
evals = sv * sv
|
||||
lambda0 = evals[0]
|
||||
cur_res["spectralnorm"] = lambda0
|
||||
cur_res["logspectralnorm"] = np.log10(lambda0)
|
||||
else:
|
||||
lambda0 = None
|
||||
if spectralnorms: # spectralnorm is the max eigenvalues
|
||||
svd = TruncatedSVD(n_components=1, n_iter=7, random_state=10)
|
||||
svd.fit(weight)
|
||||
sv = svd.singular_values_
|
||||
sv_max = np.max(sv)
|
||||
if normalize:
|
||||
evals = sv * sv / N
|
||||
else:
|
||||
evals = sv * sv
|
||||
lambda0 = evals[0]
|
||||
cur_res["spectralnorm"] = lambda0
|
||||
cur_res["logspectralnorm"] = np.log10(lambda0)
|
||||
else:
|
||||
lambda0 = None
|
||||
|
||||
if M < min_size:
|
||||
summary = "Weight matrix {}/{} ({},{}): Skipping: too small (<{})".format(i + 1, count, M, N, min_size)
|
||||
cur_res["summary"] = summary
|
||||
continue
|
||||
elif max_size > 0 and M > max_size:
|
||||
summary = "Weight matrix {}/{} ({},{}): Skipping: too big (testing) (>{})".format(i + 1, count, M, N, max_size)
|
||||
cur_res["summary"] = summary
|
||||
continue
|
||||
else:
|
||||
summary = []
|
||||
if alphas:
|
||||
import powerlaw
|
||||
svd = TruncatedSVD(n_components=M - 1, n_iter=7, random_state=10)
|
||||
svd.fit(weight.astype(float))
|
||||
sv = svd.singular_values_
|
||||
if normalize: evals = sv * sv / N
|
||||
else: evals = sv * sv
|
||||
if M < min_size:
|
||||
summary = "Weight matrix {}/{} ({},{}): Skipping: too small (<{})".format(i + 1, count, M, N, min_size)
|
||||
cur_res["summary"] = summary
|
||||
continue
|
||||
elif max_size > 0 and M > max_size:
|
||||
summary = "Weight matrix {}/{} ({},{}): Skipping: too big (testing) (>{})".format(
|
||||
i + 1, count, M, N, max_size
|
||||
)
|
||||
cur_res["summary"] = summary
|
||||
continue
|
||||
else:
|
||||
summary = []
|
||||
if alphas:
|
||||
import powerlaw
|
||||
|
||||
lambda_max = np.max(evals)
|
||||
fit = powerlaw.Fit(evals, xmax=lambda_max, verbose=False)
|
||||
alpha = fit.alpha
|
||||
cur_res["alpha"] = alpha
|
||||
D = fit.D
|
||||
cur_res["D"] = D
|
||||
cur_res["lambda_min"] = np.min(evals)
|
||||
cur_res["lambda_max"] = lambda_max
|
||||
alpha_weighted = alpha * np.log10(lambda_max)
|
||||
cur_res["alpha_weighted"] = alpha_weighted
|
||||
tolerance = lambda_max * M * np.finfo(np.max(sv)).eps
|
||||
cur_res["rank_loss"] = np.count_nonzero(sv > tolerance, axis=-1)
|
||||
svd = TruncatedSVD(n_components=M - 1, n_iter=7, random_state=10)
|
||||
svd.fit(weight.astype(float))
|
||||
sv = svd.singular_values_
|
||||
if normalize:
|
||||
evals = sv * sv / N
|
||||
else:
|
||||
evals = sv * sv
|
||||
|
||||
logpnorm = np.log10(np.sum([ev ** alpha for ev in evals]))
|
||||
cur_res["logpnorm"] = logpnorm
|
||||
lambda_max = np.max(evals)
|
||||
fit = powerlaw.Fit(evals, xmax=lambda_max, verbose=False)
|
||||
alpha = fit.alpha
|
||||
cur_res["alpha"] = alpha
|
||||
D = fit.D
|
||||
cur_res["D"] = D
|
||||
cur_res["lambda_min"] = np.min(evals)
|
||||
cur_res["lambda_max"] = lambda_max
|
||||
alpha_weighted = alpha * np.log10(lambda_max)
|
||||
cur_res["alpha_weighted"] = alpha_weighted
|
||||
tolerance = lambda_max * M * np.finfo(np.max(sv)).eps
|
||||
cur_res["rank_loss"] = np.count_nonzero(sv > tolerance, axis=-1)
|
||||
|
||||
summary.append(
|
||||
"Weight matrix {}/{} ({},{}): Alpha: {}, Alpha Weighted: {}, D: {}, pNorm {}".format(i + 1, count, M, N, alpha,
|
||||
alpha_weighted, D,
|
||||
logpnorm))
|
||||
logpnorm = np.log10(np.sum([ev ** alpha for ev in evals]))
|
||||
cur_res["logpnorm"] = logpnorm
|
||||
|
||||
if lognorms:
|
||||
norm = np.linalg.norm(weight) # Frobenius Norm
|
||||
cur_res["norm"] = norm
|
||||
lognorm = np.log10(norm)
|
||||
cur_res["lognorm"] = lognorm
|
||||
summary.append(
|
||||
"Weight matrix {}/{} ({},{}): Alpha: {}, Alpha Weighted: {}, D: {}, pNorm {}".format(
|
||||
i + 1, count, M, N, alpha, alpha_weighted, D, logpnorm
|
||||
)
|
||||
)
|
||||
|
||||
X = np.dot(weight.T, weight)
|
||||
if normalize: X = X / N
|
||||
normX = np.linalg.norm(X) # Frobenius Norm
|
||||
cur_res["normX"] = normX
|
||||
lognormX = np.log10(normX)
|
||||
cur_res["lognormX"] = lognormX
|
||||
if lognorms:
|
||||
norm = np.linalg.norm(weight) # Frobenius Norm
|
||||
cur_res["norm"] = norm
|
||||
lognorm = np.log10(norm)
|
||||
cur_res["lognorm"] = lognorm
|
||||
|
||||
summary.append(
|
||||
"Weight matrix {}/{} ({},{}): LogNorm: {} ; LogNormX: {}".format(i + 1, count, M, N, lognorm, lognormX))
|
||||
X = np.dot(weight.T, weight)
|
||||
if normalize:
|
||||
X = X / N
|
||||
normX = np.linalg.norm(X) # Frobenius Norm
|
||||
cur_res["normX"] = normX
|
||||
lognormX = np.log10(normX)
|
||||
cur_res["lognormX"] = lognormX
|
||||
|
||||
if softranks:
|
||||
softrank = norm ** 2 / sv_max ** 2
|
||||
softranklog = np.log10(softrank)
|
||||
softranklogratio = lognorm / np.log10(sv_max)
|
||||
cur_res["softrank"] = softrank
|
||||
cur_res["softranklog"] = softranklog
|
||||
cur_res["softranklogratio"] = softranklogratio
|
||||
summary += "{}. Softrank: {}. Softrank log: {}. Softrank log ratio: {}".format(summary, softrank, softranklog,
|
||||
softranklogratio)
|
||||
cur_res["summary"] = "\n".join(summary)
|
||||
return results
|
||||
summary.append(
|
||||
"Weight matrix {}/{} ({},{}): LogNorm: {} ; LogNormX: {}".format(i + 1, count, M, N, lognorm, lognormX)
|
||||
)
|
||||
|
||||
if softranks:
|
||||
softrank = norm ** 2 / sv_max ** 2
|
||||
softranklog = np.log10(softrank)
|
||||
softranklogratio = lognorm / np.log10(sv_max)
|
||||
cur_res["softrank"] = softrank
|
||||
cur_res["softranklog"] = softranklog
|
||||
cur_res["softranklogratio"] = softranklogratio
|
||||
summary += "{}. Softrank: {}. Softrank log: {}. Softrank log ratio: {}".format(
|
||||
summary, softrank, softranklog, softranklogratio
|
||||
)
|
||||
cur_res["summary"] = "\n".join(summary)
|
||||
return results
|
||||
|
||||
|
||||
def compute_details(results):
|
||||
"""
|
||||
Return a pandas data frame.
|
||||
"""
|
||||
final_summary = OrderedDict()
|
||||
"""
|
||||
Return a pandas data frame.
|
||||
"""
|
||||
final_summary = OrderedDict()
|
||||
|
||||
metrics = {
|
||||
# key in "results" : pretty print name
|
||||
"check": "Check",
|
||||
"checkTF": "CheckTF",
|
||||
"norm": "Norm",
|
||||
"lognorm": "LogNorm",
|
||||
"normX": "Norm X",
|
||||
"lognormX": "LogNorm X",
|
||||
"alpha": "Alpha",
|
||||
"alpha_weighted": "Alpha Weighted",
|
||||
"spectralnorm": "Spectral Norm",
|
||||
"logspectralnorm": "Log Spectral Norm",
|
||||
"softrank": "Softrank",
|
||||
"softranklog": "Softrank Log",
|
||||
"softranklogratio": "Softrank Log Ratio",
|
||||
"sigma_mp": "Marchenko-Pastur (MP) fit sigma",
|
||||
"numofSpikes": "Number of spikes per MP fit",
|
||||
"ratio_numofSpikes": "aka, percent_mass, Number of spikes / total number of evals",
|
||||
"softrank_mp": "Softrank for MP fit",
|
||||
"logpnorm": "alpha pNorm"
|
||||
}
|
||||
metrics = {
|
||||
# key in "results" : pretty print name
|
||||
"check": "Check",
|
||||
"checkTF": "CheckTF",
|
||||
"norm": "Norm",
|
||||
"lognorm": "LogNorm",
|
||||
"normX": "Norm X",
|
||||
"lognormX": "LogNorm X",
|
||||
"alpha": "Alpha",
|
||||
"alpha_weighted": "Alpha Weighted",
|
||||
"spectralnorm": "Spectral Norm",
|
||||
"logspectralnorm": "Log Spectral Norm",
|
||||
"softrank": "Softrank",
|
||||
"softranklog": "Softrank Log",
|
||||
"softranklogratio": "Softrank Log Ratio",
|
||||
"sigma_mp": "Marchenko-Pastur (MP) fit sigma",
|
||||
"numofSpikes": "Number of spikes per MP fit",
|
||||
"ratio_numofSpikes": "aka, percent_mass, Number of spikes / total number of evals",
|
||||
"softrank_mp": "Softrank for MP fit",
|
||||
"logpnorm": "alpha pNorm",
|
||||
}
|
||||
|
||||
metrics_stats = []
|
||||
for metric in metrics:
|
||||
metrics_stats.append("{}_min".format(metric))
|
||||
metrics_stats.append("{}_max".format(metric))
|
||||
metrics_stats.append("{}_avg".format(metric))
|
||||
|
||||
metrics_stats.append("{}_compound_min".format(metric))
|
||||
metrics_stats.append("{}_compound_max".format(metric))
|
||||
metrics_stats.append("{}_compound_avg".format(metric))
|
||||
|
||||
columns = ["layer_id", "layer_type", "N", "M", "layer_count", "slice",
|
||||
"slice_count", "level", "comment"] + [*metrics] + metrics_stats
|
||||
|
||||
metrics_values = {}
|
||||
metrics_values_compound = {}
|
||||
|
||||
for metric in metrics:
|
||||
metrics_values[metric] = []
|
||||
metrics_values_compound[metric] = []
|
||||
|
||||
layer_count = 0
|
||||
for layer_id, result in results.items():
|
||||
layer_count += 1
|
||||
|
||||
layer_type = np.NAN
|
||||
if "layer_type" in result:
|
||||
layer_type = str(result["layer_type"]).replace("LAYER_TYPE.", "")
|
||||
|
||||
compounds = {} # temp var
|
||||
metrics_stats = []
|
||||
for metric in metrics:
|
||||
compounds[metric] = []
|
||||
metrics_stats.append("{}_min".format(metric))
|
||||
metrics_stats.append("{}_max".format(metric))
|
||||
metrics_stats.append("{}_avg".format(metric))
|
||||
|
||||
slice_count, Ntotal, Mtotal = 0, 0, 0
|
||||
for slice_id, summary in result.items():
|
||||
if not str(slice_id).isdigit():
|
||||
continue
|
||||
slice_count += 1
|
||||
metrics_stats.append("{}_compound_min".format(metric))
|
||||
metrics_stats.append("{}_compound_max".format(metric))
|
||||
metrics_stats.append("{}_compound_avg".format(metric))
|
||||
|
||||
N = np.NAN
|
||||
if "N" in summary:
|
||||
N = summary["N"]
|
||||
Ntotal += N
|
||||
columns = (
|
||||
["layer_id", "layer_type", "N", "M", "layer_count", "slice", "slice_count", "level", "comment"]
|
||||
+ [*metrics]
|
||||
+ metrics_stats
|
||||
)
|
||||
|
||||
M = np.NAN
|
||||
if "M" in summary:
|
||||
M = summary["M"]
|
||||
Mtotal += M
|
||||
metrics_values = {}
|
||||
metrics_values_compound = {}
|
||||
|
||||
data = {"layer_id": layer_id, "layer_type": layer_type, "N": N, "M": M, "slice": slice_id, "level": "SLICE",
|
||||
"comment": "Slice level"}
|
||||
for metric in metrics:
|
||||
if metric in summary:
|
||||
value = summary[metric]
|
||||
if value is not None:
|
||||
metrics_values[metric].append(value)
|
||||
compounds[metric].append(value)
|
||||
data[metric] = value
|
||||
for metric in metrics:
|
||||
metrics_values[metric] = []
|
||||
metrics_values_compound[metric] = []
|
||||
|
||||
data = {"layer_id": layer_id, "layer_type": layer_type, "N": Ntotal, "M": Mtotal, "slice_count": slice_count,
|
||||
"level": "LAYER", "comment": "Layer level"}
|
||||
# Compute the compound value over the slices
|
||||
for metric, value in compounds.items():
|
||||
count = len(value)
|
||||
if count == 0:
|
||||
continue
|
||||
layer_count = 0
|
||||
for layer_id, result in results.items():
|
||||
layer_count += 1
|
||||
|
||||
compound = np.mean(value)
|
||||
metrics_values_compound[metric].append(compound)
|
||||
data[metric] = compound
|
||||
layer_type = np.NAN
|
||||
if "layer_type" in result:
|
||||
layer_type = str(result["layer_type"]).replace("LAYER_TYPE.", "")
|
||||
|
||||
data = {"layer_count": layer_count, "level": "NETWORK", "comment": "Network Level"}
|
||||
for metric, metric_name in metrics.items():
|
||||
if metric not in metrics_values or len(metrics_values[metric]) == 0:
|
||||
continue
|
||||
compounds = {} # temp var
|
||||
for metric in metrics:
|
||||
compounds[metric] = []
|
||||
|
||||
values = metrics_values[metric]
|
||||
minimum = min(values)
|
||||
maximum = max(values)
|
||||
avg = np.mean(values)
|
||||
final_summary[metric] = avg
|
||||
# print("{}: min: {}, max: {}, avg: {}".format(metric_name, minimum, maximum, avg))
|
||||
data["{}_min".format(metric)] = minimum
|
||||
data["{}_max".format(metric)] = maximum
|
||||
data["{}_avg".format(metric)] = avg
|
||||
slice_count, Ntotal, Mtotal = 0, 0, 0
|
||||
for slice_id, summary in result.items():
|
||||
if not str(slice_id).isdigit():
|
||||
continue
|
||||
slice_count += 1
|
||||
|
||||
values = metrics_values_compound[metric]
|
||||
minimum = min(values)
|
||||
maximum = max(values)
|
||||
avg = np.mean(values)
|
||||
final_summary["{}_compound".format(metric)] = avg
|
||||
# print("{} compound: min: {}, max: {}, avg: {}".format(metric_name, minimum, maximum, avg))
|
||||
data["{}_compound_min".format(metric)] = minimum
|
||||
data["{}_compound_max".format(metric)] = maximum
|
||||
data["{}_compound_avg".format(metric)] = avg
|
||||
N = np.NAN
|
||||
if "N" in summary:
|
||||
N = summary["N"]
|
||||
Ntotal += N
|
||||
|
||||
return final_summary
|
||||
M = np.NAN
|
||||
if "M" in summary:
|
||||
M = summary["M"]
|
||||
Mtotal += M
|
||||
|
||||
data = {
|
||||
"layer_id": layer_id,
|
||||
"layer_type": layer_type,
|
||||
"N": N,
|
||||
"M": M,
|
||||
"slice": slice_id,
|
||||
"level": "SLICE",
|
||||
"comment": "Slice level",
|
||||
}
|
||||
for metric in metrics:
|
||||
if metric in summary:
|
||||
value = summary[metric]
|
||||
if value is not None:
|
||||
metrics_values[metric].append(value)
|
||||
compounds[metric].append(value)
|
||||
data[metric] = value
|
||||
|
||||
data = {
|
||||
"layer_id": layer_id,
|
||||
"layer_type": layer_type,
|
||||
"N": Ntotal,
|
||||
"M": Mtotal,
|
||||
"slice_count": slice_count,
|
||||
"level": "LAYER",
|
||||
"comment": "Layer level",
|
||||
}
|
||||
# Compute the compound value over the slices
|
||||
for metric, value in compounds.items():
|
||||
count = len(value)
|
||||
if count == 0:
|
||||
continue
|
||||
|
||||
compound = np.mean(value)
|
||||
metrics_values_compound[metric].append(compound)
|
||||
data[metric] = compound
|
||||
|
||||
data = {"layer_count": layer_count, "level": "NETWORK", "comment": "Network Level"}
|
||||
for metric, metric_name in metrics.items():
|
||||
if metric not in metrics_values or len(metrics_values[metric]) == 0:
|
||||
continue
|
||||
|
||||
values = metrics_values[metric]
|
||||
minimum = min(values)
|
||||
maximum = max(values)
|
||||
avg = np.mean(values)
|
||||
final_summary[metric] = avg
|
||||
# print("{}: min: {}, max: {}, avg: {}".format(metric_name, minimum, maximum, avg))
|
||||
data["{}_min".format(metric)] = minimum
|
||||
data["{}_max".format(metric)] = maximum
|
||||
data["{}_avg".format(metric)] = avg
|
||||
|
||||
values = metrics_values_compound[metric]
|
||||
minimum = min(values)
|
||||
maximum = max(values)
|
||||
avg = np.mean(values)
|
||||
final_summary["{}_compound".format(metric)] = avg
|
||||
# print("{} compound: min: {}, max: {}, avg: {}".format(metric_name, minimum, maximum, avg))
|
||||
data["{}_compound_min".format(metric)] = minimum
|
||||
data["{}_compound_max".format(metric)] = maximum
|
||||
data["{}_compound_avg".format(metric)] = avg
|
||||
|
||||
return final_summary
|
||||
|
||||
|
||||
def analyze(model: nn.Module, min_size=50, max_size=0,
|
||||
alphas: bool = False, lognorms: bool = True, spectralnorms: bool = False,
|
||||
softranks: bool = False, normalize: bool = False, glorot_fix: bool = False):
|
||||
"""
|
||||
Analyze the weight matrices of a model.
|
||||
:param model: A PyTorch model
|
||||
:param min_size: The minimum weight matrix size to analyze.
|
||||
:param max_size: The maximum weight matrix size to analyze (0 = no limit).
|
||||
:param alphas: Compute the power laws (alpha) of the weight matrices.
|
||||
Time consuming so disabled by default (use lognorm if you want speed)
|
||||
:param lognorms: Compute the log norms of the weight matrices.
|
||||
:param spectralnorms: Compute the spectral norm (max eigenvalue) of the weight matrices.
|
||||
:param softranks: Compute the soft norm (i.e. StableRank) of the weight matrices.
|
||||
:param normalize: Normalize or not.
|
||||
:param glorot_fix:
|
||||
:return: (a dict of all layers' results, a dict of the summarized info)
|
||||
"""
|
||||
names, modules = [], []
|
||||
for name, module in model.named_modules():
|
||||
if isinstance(module, available_module_types()):
|
||||
names.append(name)
|
||||
modules.append(module)
|
||||
# print('There are {:} layers to be analyzed in this model.'.format(len(modules)))
|
||||
all_results = OrderedDict()
|
||||
for index, module in enumerate(modules):
|
||||
if isinstance(module, nn.Linear):
|
||||
weights = [module.weight.cpu().detach().numpy()]
|
||||
else:
|
||||
weights = get_conv2D_Wmats(module.weight.cpu().detach().numpy())
|
||||
results = analyze_weights(weights, min_size, max_size, alphas, lognorms, spectralnorms, softranks, normalize, glorot_fix)
|
||||
results['id'] = index
|
||||
results['type'] = type(module)
|
||||
all_results[index] = results
|
||||
summary = compute_details(all_results)
|
||||
return all_results, summary
|
||||
def analyze(
|
||||
model: nn.Module,
|
||||
min_size=50,
|
||||
max_size=0,
|
||||
alphas: bool = False,
|
||||
lognorms: bool = True,
|
||||
spectralnorms: bool = False,
|
||||
softranks: bool = False,
|
||||
normalize: bool = False,
|
||||
glorot_fix: bool = False,
|
||||
):
|
||||
"""
|
||||
Analyze the weight matrices of a model.
|
||||
:param model: A PyTorch model
|
||||
:param min_size: The minimum weight matrix size to analyze.
|
||||
:param max_size: The maximum weight matrix size to analyze (0 = no limit).
|
||||
:param alphas: Compute the power laws (alpha) of the weight matrices.
|
||||
Time consuming so disabled by default (use lognorm if you want speed)
|
||||
:param lognorms: Compute the log norms of the weight matrices.
|
||||
:param spectralnorms: Compute the spectral norm (max eigenvalue) of the weight matrices.
|
||||
:param softranks: Compute the soft norm (i.e. StableRank) of the weight matrices.
|
||||
:param normalize: Normalize or not.
|
||||
:param glorot_fix:
|
||||
:return: (a dict of all layers' results, a dict of the summarized info)
|
||||
"""
|
||||
names, modules = [], []
|
||||
for name, module in model.named_modules():
|
||||
if isinstance(module, available_module_types()):
|
||||
names.append(name)
|
||||
modules.append(module)
|
||||
# print('There are {:} layers to be analyzed in this model.'.format(len(modules)))
|
||||
all_results = OrderedDict()
|
||||
for index, module in enumerate(modules):
|
||||
if isinstance(module, nn.Linear):
|
||||
weights = [module.weight.cpu().detach().numpy()]
|
||||
else:
|
||||
weights = get_conv2D_Wmats(module.weight.cpu().detach().numpy())
|
||||
results = analyze_weights(
|
||||
weights, min_size, max_size, alphas, lognorms, spectralnorms, softranks, normalize, glorot_fix
|
||||
)
|
||||
results["id"] = index
|
||||
results["type"] = type(module)
|
||||
all_results[index] = results
|
||||
summary = compute_details(all_results)
|
||||
return all_results, summary
|
||||
|
Reference in New Issue
Block a user