pytorch检查模型各部分输入与输出

最新推荐文章于 2026-05-14 05:10:00 发布

原创最新推荐文章于 2026-05-14 05:10:00 发布 · 3.7k 阅读

2 ·

本内容遵循CC 4.0 BY-SA版权协议

标签

#pytorch #深度学习 #机器学习

收录于

自学专栏收录该内容

17 篇文章

订阅专栏

本文介绍了一个用于手写数字识别的卷积神经网络（CNN）模型实现过程，该模型包含多个卷积层、池化层及全连接层，并通过逐步展示输入数据经过各层后的形状变化来验证模型的有效性。

import torch
from torch import nn
from d2l import torch as d2l

class Reshape(torch.nn.Module):
    def forward(selfself,x):
        return x.view(-1,1,28,28)

net=torch.nn.Sequential(Reshape(),
                        nn.Conv2d(1,6,kernel_size=5,padding=2),
                        nn.Sigmoid(),
                        nn.AvgPool2d(2,
                                     stride=2),
                        nn.Conv2d(6,16,kernel_size=5),nn.Sigmoid(),
                        nn.AvgPool2d(kernel_size=2,stride=2),
                        nn.Flatten(),
                        nn.Linear(16*5*5,120),nn.Sigmoid(),
                        nn.Linear(120,84),nn.Sigmoid(),
                        nn.Linear(84,10))
# 检查模型
X=torch.rand(size=(1,1,28,28),dtype=torch.float32)
for layer in net:
    X=layer(X)
    print(layer.__class__.__name__,'output shape:\t',X.shape)