![](https://img-blog.csdnimg.cn/20201014180756918.png?x-oss-process=image/resize,m_fixed,h_64,w_64)
hook
那里,春暖花开
这个作者很懒,什么都没留下…
展开
-
手动实现计算模型的参数量和计算量
此代码来源于Higher HRNet import os import torch import torch.nn as nn import torchvision.models as models def get_model_summary(model, *input_tensors, item_length=26, verbose=True): """ :param model: :param input_tensors: :param item_length:原创 2020-06-16 20:11:12 · 649 阅读 · 0 评论 -
register_forward_hook(hook)
建议阅读官方文档说明 该函数对模块注册一个钩子函数。每次模块的forward()函数计算输出后,都会调用该钩子函数。钩子函数必须具有以下函数签名hook(module, input, output) 一个简单的conv模块调用一个计算参数量的钩子函数如下 def hook(module,input,output): class_name = str(module.__class__.__name__) if class_name.find("Conv") != -1 or class_name.原创 2020-06-16 16:53:25 · 1127 阅读 · 0 评论