当我们使用pytorch进行模型训练或测试时,有时候希望能知道模型每一层分别是什么,具有怎样的参数。此时我们可以将模型打印出来,输出每一层的名字、类型、参数等。
常用的命令行打印模型结构的方法有两种:
import torch
from torch import nn
#from d2l import torch as d2l
net1D = nn.Sequential(nn.Conv1d(1,6,kernel_size=5,padding=2),nn.Sigmoid(),
nn.AvgPool1d(kernel_size=2,stride=2),
nn.Conv1d(6,16,kernel_size=5),nn.Sigmoid(),
nn.AvgPool1d(kernel_size=2,stride=2),
nn.Flatten(),
nn.Linear(16*5,120),nn.Sigmoid(), # 这边要根据上面的输出重新计算拉平后的大小
nn.Linear(120,84),nn.Sigmoid(),
nn.Linear(84,2)
)
Y=torch.rand(size=(1,1,28),dtype=torch.float32) # 批次大小,通道数,长度
for layer in net1D:
Y=layer(Y)
print(layer.__class__.__name__, 'output shape: \t',Y.shape)
Conv1d output shape: torch.Size([1, 6, 28])
Sigmoid output shape: torch.Size([1, 6, 28])
AvgPool1d output shape: torch.Size([1, 6, 14])
Conv1d output shape: torch.Size([1, 16, 10])
Sigmoid output shape: torch.Size([1, 16, 10])
AvgPool1d output shape: torch.Size([1, 16, 5])
Flatten output shape: torch.Size([1, 80])
Linear output shape: torch.Size([1, 120])
Sigmoid output shape: torch.Size([1, 120])
Linear output shape: torch.Size([1, 84])
Sigmoid output shape: torch.Size([1, 84])
Linear output shape: torch.Size([1, 2])