公众号关注 “ML_NLP
设为 “星标”,重磅干货,第一时间送达!
作者:Big fish
地址:https://zhuanlan.zhihu.com/p/33992733
从网上各种资料加上自己实践的可用工具。

主要包括:
模型层数:print_layers_num
模型参数总量:print_model_parm_nums
模型的计算图:def print_autograd_graph():或者参见tensorboad
模型滤波器可视化:show_save_tensor
模型在具体的输入下的尺寸信息summary以及参数量:show_summary
模型计算量:print_model_parm_flops
格式较混乱,但上述代码均可用,后续会继续整理。
#coding:utf8import torchimport torchvisionimport torch.nn as nnfrom torch.autograd import Variableimport torchvision.models as modelsimport numpy as npdeftest(): model = models.resnet18()print model.layer1[0].conv1.weight.dataprint model.layer1[0].conv1.__class__#<class 'torch.nn.modules.conv.Conv2d'>print model.layer1[0].conv1.kernel_size input = torch.autograd.Variable(torch.randn(20, 16, 50, 100))print input.size()print np.prod(input.size())defprint_model_parm_nums(): model = models.alexnet() total = sum([param.nelement() for param in model.parameters()]) print(' + Number of params: %.2fM' % (total / 1e6))defprint_model_parm_flops():# prods = {}# def save_prods(self, input, output):# print 'flops:{}'.format(self.__class__.__name__)# print 'input:{}'.format(input)# print '_dim:{}'.format(input[0].dim())# print 'input_shape:{}'.format(np.prod(input[0].shape))# grads.append(np.prod(input[0].shape)) prods = {}defsave_hook(name):defhook_per(self, input, output):# print 'flops:{}'.format(self.__class__.__name__)# print 'input:{}'.format(input)# print '_dim:{}'.format(input[0].dim())# print 'input_shape:{}'.format(np.prod(input[0].shape))# prods.append(np.prod(input[0].shape)) prods[name] = np.prod(input[0].shape)# prods.append(np.prod(input[0].shape))return hook_per list_1=[]defsimple_hook(self, input, output): list_1.append(np.prod(input[0].shape)) list_2={}defsimple_hook2(self, input, output): list_2['names'] = np.prod(input[0].shape) multiply_adds = False list_conv=[]defconv_hook(self, input, output): batch_size, input_channels, input_height, input_width = input[0].size() output_channels, output_height, output_width = output[0].size() kernel_ops = self.kernel_size[0] * self.kernel_size[1] * (self.in_channels / self.groups) * (2if multiply_adds else1) bias_ops = 1if self.bias isnotNoneelse0 params = output_channels * (kernel_ops + bias_ops) flops = batch_size * params * output_height * output_width list_conv.append(flops) list_linear=[] deflinear_hook(self, input, output): batch_size = input[0].size(0) if input[0].dim() == 2else1 weight_ops = self.weight.nelement() * (2if multiply_adds else1) bias_ops = self.bias.nelement() flops = batch_size * (weight_ops + bias_ops) list_linear.append(flops) list_bn=[] defbn_hook(self, input, output): list_bn.append(input[0].nelement()) list_relu=[] defrelu_hook(self, input, output): list_relu.append(input[0].nelement()) list_pooling=[]defpooling_hook(self, input, output): batch_size, input_channels, input_height, input_width = input[0].size() output_channels, output_height, output_width = output[0].size() kernel_ops = self.kernel_size * self.kernel_size bias_ops = 0 params = output_channels * (kernel_ops + bias_ops) flops = batch_size * params * output_height * output_width list_pooling.append(flops)deffoo(net): childrens = list(net.children())ifnot childrens:if isinstance(net, torch.nn.Conv2d):# net.register_forward_hook(save_hook(net.__class__.__name__))# net.register_forward_hook(simple_hook)# net.register_forward_hook(simple_hook2) net.register_forward_hook(conv_hook)if isinstance(net, torch.nn.Linear): net.register_forward_hook(linear_hook)if isinstance(net, torch.nn.BatchNorm2d): net.register_forward_hook(bn_hook)if isinstance(net, torch.nn.ReLU): net.register_forward_hook(relu_hook)if isinstance(net, torch.nn.MaxPool2d) or isinstance(net, torch.nn.AvgPool2d): net.register_forward_hook(pooling_hook)returnfor c in childrens: foo(c) resnet = models.alexnet() foo(resnet) input = Variable(torch.rand(3,224,224).unsqueeze(0), requires_grad = True) out = resnet(input) total_flops = (sum(list_conv) + sum(list_linear) + sum(list_bn) + sum(list_relu) + sum(list_pooling)) print(' + Number of FLOPs: %.2fG' % (total_flops / 1e9))# print list_bn# print 'prods:{}'.format(prods)# print 'list_1:{}'.format(list_1)# print 'list_2:{}'.format(list_2)# print 'list_final:{}'.format(list_final)defprint_forward(): model = torchvision.models.resnet18() select_layer = model.layer1[0].conv1 grads={}defsave_grad(name):defhook(self, input, output): grads[name] = inputreturn hook select_layer.register_forward_hook(save_grad('select_layer')) input = Variable(torch.rand(3,224,224).unsqueeze(0), requires_grad = True) out = model(input)# print grads['select_layer']print gradsdefprint_value(): grads = {}defsave_grad(name):defhook(grad): grads[name] = gradreturn hook x = Variable(torch.randn(1,1), requires_grad=True) y = 3*x z = y**2# In here, save_grad('y') returns a hook (a function) that keeps 'y' as name y.register_hook(save_grad('y')) z.register_hook(save_grad('z')) z.backward()print'HW' print("grads['y']: {}".format(grads['y'])) print(grads['z'])defprint_layers_num():# resnet = models.resnet18() resnet = models.resnet18()deffoo(net): childrens = list(net.children())ifnot childrens:if isinstance(net, torch.nn.Conv2d):print' '#可以用来统计不同层的个数# net.register_backward_hook(print)return1 count = 0for c in childrens: count += foo(c)return count print(foo(resnet))defcheck_summary():deftorch_summarize(model, show_weights=True, show_parameters=True):"""Summarizes torch model by showing trainable parameters and weights."""from torch.nn.modules.module import _addindent tmpstr = model.__class__.__name__ + ' (\n'for key, module in model._modules.items():# if it contains layers let call it recursively to get params and weightsif type(module) in [ torch.nn.modules.container.Container, torch.nn.modules.container.Sequential ]: modstr = torch_summarize(module)else: modstr = module.__repr__() modstr = _addindent(modstr, 2) params = sum([np.prod(p.size()) for p in module.parameters()]) weights = tuple([tuple(p.size()) for p in module.parameters()]) tmpstr += ' (' + key + '): ' + modstr if show_weights: tmpstr += ', weights={}'.format(weights)if show_parameters: tmpstr += ', parameters={}'.format(params) tmpstr += '\n' tmpstr = tmpstr + ')'return tmpstr# Testimport torchvision.models as models model = models.alexnet() print(torch_summarize(model))#https://gist.github.com/wassname/0fb8f95e4272e6bdd27bd7df386716b7#summarize a torch model like in keras, showing parameters and output shapedefshow_summary():from collections import OrderedDictimport pandas as pdimport numpy as npimport torchfrom torch.autograd import Variableimport torch.nn.functional as Ffrom torch import nndefget_names_dict(model):""" Recursive walk to get names including path """ names = {}def_get_names(module, parent_name=''):for key, module in module.named_children(): name = parent_name + '.' + key if parent_name else key names[name]=moduleif isinstance(module, torch.nn.Module): _get_names(module, parent_name=name) _get_names(model)return namesdeftorch_summarize_df(input_size, model, weights=False, input_shape=True, nb_trainable=False):""" Summarizes torch model by showing trainable parameters and weights. author: wassname url: https://gist.github.com/wassname/0fb8f95e4272e6bdd27bd7df386716b7 license: MIT Modified from: - https://github.com/pytorch/pytorch/issues/2001#issuecomment-313735757 - https://gist.github.com/wassname/0fb8f95e4272e6bdd27bd7df386716b7/ Usage: import torchvision.models as models model = models.alexnet() df = torch_summarize_df(input_size=(3, 224,224), model=model) print(df) # name class_name input_shape output_shape nb_params # 1 features=>0 Conv2d (-1, 3, 224, 224) (-1, 64, 55, 55) 23296#(3*11*11+1)*64 # 2 features=>1 ReLU (-1, 64, 55, 55) (-1, 64, 55, 55) 0 # ... """defregister_hook(module):defhook(module, input, output): name = ''for key, item in names.items():if item == module: name = key#<class 'torch.nn.modules.conv.Conv2d'> class_name = str(module.__class__).split('.')[-1].split("'")[0] module_idx = len(summary) m_key = module_idx + 1 summary[m_key] = OrderedDict() summary[m_key]['name'] = name summary[m_key]['class_name'] = class_nameif input_shape: summary[m_key]['input_shape'] = (-1, ) + tuple(input[0].size())[1:] summary[m_key]['output_shape'] = (-1, ) + tuple(output.size())[1:]if weights: summary[m_key]['weights'] = list( [tuple(p.size()) for p in module.parameters()])# summary[m_key]['trainable'] = any([p.requires_grad for p in module.parameters()])if nb_trainable: params_trainable = sum([torch.LongTensor(list(p.size())).prod() for p in module.parameters() if p.requires_grad]) summary[m_key]['nb_trainable'] = params_trainable params = sum([torch.LongTensor(list(p.size())).prod() for p in module.parameters()]) summary[m_key]['nb_params'] = paramsifnot isinstance(module, nn.Sequential) and \not isinstance(module, nn.ModuleList) and \not (module == model): hooks.append(module.register_forward_hook(hook))# Names are stored in parent and path+name is unique not the name names = get_names_dict(model)# check if there are multiple inputs to the networkif isinstance(input_size[0], (list, tuple)): x = [Variable(torch.rand(1, *in_size)) for in_size in input_size]else: x = Variable(torch.rand(1, *input_size))if next(model.parameters()).is_cuda: x = x.cuda()# create properties summary = OrderedDict() hooks = []# register hook model.apply(register_hook)# make a forward pass model(x)# remove these hooksfor h in hooks: h.remove()# make dataframe df_summary = pd.DataFrame.from_dict(summary, orient='index')return df_summary# Test on alexnetimport torchvision.models as models model = models.alexnet() df = torch_summarize_df(input_size=(3, 224, 224), model=model) print(df)# # Output# name class_name input_shape output_shape nb_params# 1 features=>0 Conv2d (-1, 3, 224, 224) (-1, 64, 55, 55) 23296#nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),# 2 features=>1 ReLU (-1, 64, 55, 55) (-1, 64, 55, 55) 0# 3 features=>2 MaxPool2d (-1, 64, 55, 55) (-1, 64, 27, 27) 0# 4 features=>3 Conv2d (-1, 64, 27, 27) (-1, 192, 27, 27) 307392# 5 features=>4 ReLU (-1, 192, 27, 27) (-1, 192, 27, 27) 0# 6 features=>5 MaxPool2d (-1, 192, 27, 27) (-1, 192, 13, 13) 0# 7 features=>6 Conv2d (-1, 192, 13, 13) (-1, 384, 13, 13) 663936# 8 features=>7 ReLU (-1, 384, 13, 13) (-1, 384, 13, 13) 0# 9 features=>8 Conv2d (-1, 384, 13, 13) (-1, 256, 13, 13) 884992# 10 features=>9 ReLU (-1, 256, 13, 13) (-1, 256, 13, 13) 0# 11 features=>10 Conv2d (-1, 256, 13, 13) (-1, 256, 13, 13) 590080# 12 features=>11 ReLU (-1, 256, 13, 13) (-1, 256, 13, 13) 0# 13 features=>12 MaxPool2d (-1, 256, 13, 13) (-1, 256, 6, 6) 0# 14 classifier=>0 Dropout (-1, 9216) (-1, 9216) 0# 15 classifier=>1 Linear (-1, 9216) (-1, 4096) 37752832# 16 classifier=>2 ReLU (-1, 4096) (-1, 4096) 0# 17 classifier=>3 Dropout (-1, 4096) (-1, 4096) 0# 18 classifier=>4 Linear (-1, 4096) (-1, 4096) 16781312# 19 classifier=>5 ReLU (-1, 4096) (-1, 4096) 0# 20 classifier=>6 Linear (-1, 4096) (-1, 1000) 4097000defshow_save_tensor():import torchfrom torchvision import utilsimport torchvision.models as modelsfrom matplotlib import pyplot as pltdefvis_tensor(tensor, ch = 0, all_kernels=False, nrow=8, padding = 2):''' ch: channel for visualization allkernels: all kernels for visualization ''' n,c,h,w = tensor.shapeif all_kernels: tensor = tensor.view(n*c ,-1, w, h)elif c != 3: tensor = tensor[:, ch,:,:].unsqueeze(dim=1) rows = np.min((tensor.shape[0]//nrow + 1, 64 )) grid = utils.make_grid(tensor, nrow=nrow, normalize=True, padding=padding)# plt.figure(figsize=(nrow,rows)) plt.imshow(grid.numpy().transpose((1, 2, 0)))#CHW HWCdefsave_tensor(tensor, filename, ch=0, all_kernels=False, nrow=8, padding=2): n,c,h,w = tensor.shapeif all_kernels: tensor = tensor.view(n*c ,-1, w, h)elif c != 3: tensor = tensor[:, ch,:,:].unsqueeze(dim=1) utils.save_image(tensor, filename, nrow = nrow,normalize=True, padding=padding) vgg = models.resnet18(pretrained=True) mm = vgg.double() filters = mm.modules body_model = [i for i in mm.children()][0]# layer1 = body_model[0] layer1 = body_model tensor = layer1.weight.data.clone() vis_tensor(tensor) save_tensor(tensor,'test.png') plt.axis('off') plt.ioff() plt.show()defprint_autograd_graph():from graphviz import Digraphimport torchfrom torch.autograd import Variabledefmake_dot(var, params=None):""" Produces Graphviz representation of PyTorch autograd graph Blue nodes are the Variables that require grad, orange are Tensors saved for backward in torch.autograd.Function Args: var: output Variable params: dict of (name, Variable) to add names to node that require grad (TODO: make optional) """if params isnotNone:#assert all(isinstance(p, Variable) for p in params.values()) param_map = {id(v): k for k, v in params.items()} node_attr = dict(style='filled', shape='box', align='left', fontsize='12', ranksep='0.1', height='0.2') dot = Digraph(node_attr=node_attr, graph_attr=dict(size="12,12")) seen = set()defsize_to_str(size):return'('+(', ').join(['%d' % v for v in size])+')'defadd_nodes(var):if var notin seen:if torch.is_tensor(var): dot.node(str(id(var)), size_to_str(var.size()), fillcolor='orange')elif hasattr(var, 'variable'): u = var.variable#name = param_map[id(u)] if params is not None else ''#node_name = '%s\n %s' % (name, size_to_str(u.size())) node_name = '%s\n %s' % (param_map.get(id(u.data)), size_to_str(u.size())) dot.node(str(id(var)), node_name, fillcolor='lightblue')else: dot.node(str(id(var)), str(type(var).__name__)) seen.add(var)if hasattr(var, 'next_functions'):for u in var.next_functions:if u[0] isnotNone: dot.edge(str(id(u[0])), str(id(var))) add_nodes(u[0])if hasattr(var, 'saved_tensors'):for t in var.saved_tensors: dot.edge(str(id(t)), str(id(var))) add_nodes(t) add_nodes(var.grad_fn)return dotfrom torchvision import models torch.manual_seed(1) inputs = torch.randn(1,3,224,224) model = models.resnet18(pretrained=False) y = model(Variable(inputs))#print(y) g = make_dot(y, params=model.state_dict()) g.view()#gif __name__=='__main__':import fire fire. Fire()

推荐阅读:

继续阅读