PyTorch⑩---卷积神经网络_一个小的神经网络搭建

一个小的神经网络搭建

import torch
from torch import nn
from torch.nn import Conv2d, MaxPool2d, Flatten, Linear


class Demo(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1=Conv2d(in_channels=3,out_channels=32,kernel_size=5,stride=1,padding=2,dilation=1,)
        self.maxpool1=MaxPool2d(kernel_size=2,)
        self.conv2=Conv2d(in_channels=32,out_channels=32,kernel_size=5,stride=1,padding=2,)
        self.maxpool2=MaxPool2d(kernel_size=2)
        self.conv3=Conv2d(32,64,5,1,2)
        self.maxpool3=MaxPool2d(2)
        self.flatten=Flatten()#如果在这里不知道展开后是多少，直接运行到这里查看一下情况
        self.linear1=Linear(1024,64)
        #为什么最后一个线性层是10，因为label是10，根据概率进行分类
        self.linear2=Linear(64,10)
    def forward(self,x):
        x=self.conv1(x)
        x=self.maxpool1(x)
        x=self.conv2(x)
        x=self.maxpool2(x)
        x=self.conv3(x)
        x=self.maxpool3(x)
        x=self.flatten(x)
        x=self.linear1(x)
        x=self.linear2(x)
        return x

demo=Demo()
print(demo)
'''对网络步骤进行验证'''
input=torch.ones((64,3,32,32))
output=demo(input)
print(output)
print(output.shape)

结果：