hook
那里,春暖花开
这个作者很懒,什么都没留下…
展开
-
手动实现计算模型的参数量和计算量
此代码来源于Higher HRNetimport osimport torchimport torch.nn as nnimport torchvision.models as modelsdef 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 · 619 阅读 · 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 · 1114 阅读 · 0 评论