Pytorch 工具之 torchsummary 及其修改

最新推荐文章于 2024-07-08 15:57:37 发布

「已注销」

最新推荐文章于 2024-07-08 15:57:37 发布

阅读量9k

点赞数 9

分类专栏： Python 文章标签： python pytorch

本文链接：https://blog.csdn.net/qq_24084387/article/details/104432604

版权

本文介绍了在超分辨率任务(SISR)中使用torchsummary进行模型参数和内存评估的经验。作者发现torchsummary的输出与实际需求不符，主要问题包括层结构层次不明显、无法清晰展示模块差异以及Forward/backward pass size计算错误。通过修改torchsummary源代码，解决了这些问题，改进包括对层的缩进展示、明确模块输入输出和修正参数大小计算。经过修改后的torchsummary能更好地适应自定义Resblocks结构的模型。

摘要由CSDN通过智能技术生成

最近在做SISR的时候需要统计模型的参数，上网搜寻了一番发现torchsummary可以对模型的参数和内存进行评估，于是尝试使用，但在使用过程中发现torchsummary的输出和自己的code对应不上，所有在其基础上进行了修改做了本土化。在此对torchsummary和自己的修改过程做记录。

torchsummary

torchsummary 的项目地址是：https://github.com/sksq96/pytorch-summary
torchsummary 的安装可以使用

pip intall torchsummary

或者

git clone https://github.com/sksq96/pytorch-summary

使用方式为：

from torchsummary import summary

summary(model, input_size, batch_size=-1, device="cuda")

以 VGG16为例，其效果为（copy自项目地址）：

import torch
from torchvision import models
from torchsummary import summary

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
vgg = models.vgg16().to(device)

summary(vgg, (3, 224, 224))

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 64, 224, 224]           1,792
              ReLU-2         [-1, 64, 224, 224]               0
            Conv2d-3         [-1, 64, 224, 224]          36,928
              ReLU-4         [-1, 64, 224, 224]               0
         MaxPool2d-5         [-1, 64, 112, 112]               0
            Conv2d-6        [-1, 128, 112, 112]          73,856
              ReLU-7        [-1, 128, 112, 112]               0
            Conv2d-8        [-1, 128, 112, 112]         147,584
              ReLU-9        [-1, 128, 112, 112]               0
        MaxPool2d-10          [-1, 128, 56, 56]               0
           Conv2d-11          [-1, 256, 56, 56]         295,168
             ReLU-12          [-1, 256, 56, 56]               0
           Conv2d-13          [-1, 256, 56, 56]         590,080
             ReLU-14          [-1, 256, 56, 56]               0
           Conv2d-15          [-1, 256, 56, 56]         590,080
             ReLU-16          [-1, 256, 56, 56]               0
        MaxPool2d-17          [-1, 256, 28, 28]               0
           Conv2d-18          [-1, 512, 28, 28]       1,180,160
             ReLU-19          [-1, 512, 28, 28]               0
           Conv2d-20          [-1, 512, 28, 28]       2,359,808
             ReLU-21          [-1, 512, 28, 28]               0
           Conv2d-22          [-1, 512, 28, 28]       2,359,808
             ReLU-23          [-1, 512, 28, 28]               0
        MaxPool2d-24          [-1, 512, 14, 14]               0
           Conv2d-25          [-1, 512, 14, 14]       2,359,808
             ReLU-26          [-1, 512, 14, 14]               0
           Conv2d-27