pytorch模型可视化的方法总结，（参数量内存，特征图内存），FLOPs和Parameters

原创

已于 2023-06-26 14:00:55 修改 · 4.6k 阅读

收录于

当前文章被以下社区和专栏收录：

于 2022-09-09 17:08:28 首次发布

本文介绍了如何使用torchsummary、torchviz、netron、tensorwatch和get_model_complexity_info等工具对PyTorch的SRCNN模型进行网络结构可视化和参数计算。通过实例展示了各方法的使用步骤，包括模型的参数量、FLOPs计算以及内存占用分析，帮助理解模型复杂性和资源需求。

文章目录

这里主要介绍pytorch 模型的网络结构的可视化
以 SRCNN 为例子来说明可视化的方法，以及参数量的计算

模型所占内存 = （参数量内存，特征图内存），
模型计算量 = （浮点数计算量）

1. torchsummary

torchinfo

class SRCNN(nn.Module):
    def __init__(self, num_channels=1):
        super(SRCNN, self).__init__()
        self.conv1 = nn.Conv2d(num_channels, 64, kernel_size=9, padding=9 // 2)
        self.conv2 = nn.Conv2d(64, 32, kernel_size=5, padding=5 // 2)
        self.conv3 = nn.Conv2d(32, num_channels, kernel_size=5, padding=5 // 2)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.conv3(x)
        return x


from torchinfo import summary
if __name__ == "__main__":
    modelviz = SRCNN()
    # 打印模型结构
    print(modelviz)
    summary(modelviz, input_size=(8, 1, 8, 8), col_names=["kernel_size", "output_size", "num_params", "mult_adds"])
    for p in modelviz.parameters():
        if p.requires_grad