Pytorch模型如何查看每层输入维度输出维度

在 PyTorch 中,可以使用 torchsummary 库来实现对 PyTorch 模型的结构及参数统计的输出,其可以方便我们查看每层输入、输出的维度以及参数数量等信息。

安装 torchsummary 库:

pip install torchsummary

使用方法如下:

import torch
from torchsummary import summary

# 定义一个模型
model = torch.nn.Sequential(
)

# 打印模型所有层的参数统计
summary(model, (3, 32, 32))

其中,model 是需要查看的模型,(3, 32, 32) 表示模型的输入维度,即 C = 3,H = 32,W = 32。运行后可以看到所有层输入、输出的维度以及参数数量等信息。

有个大坑
AttributeError: ‘xxx’ object has no attribute ‘size’
AttributeError: ‘int’ object has no attribute ‘size’
类似错误

pip uninstall torchsummary
pip install torch-summary

解决问题

猜你喜欢

转载自blog.csdn.net/qq_22815083/article/details/129622393