Pytorch入门学习写一个简单的分类程序

参考视频:https://www.bilibili.com/video/BV1AP4y167bX/?spm_id_from=333.999.0.0&vd_source=98d31d5c9db8c0021988f2c2c25a9620

网络初学

import torch
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1=nn.Conv2d(1,6,5)
        self.conv2=nn.Conv2d(6,16,5)
        self.fc1=nn.Linear(16*5*5,120)
        self.fc2=nn.Linear(120,84)
        self.fc3=nn.Linear(84,10)
    def forward(self,x):
        #x : tensor[batch,channel,H,W]
        #假设这里的x的维度是(1,1,32,32)
        x=self.conv1(x)#x的维度是(1,6,28,28)
        x=F.relu(x)#x的维度是(1,6,28,28)
        x=F.max_pool2d(x,(2,2))#x的维度是(1,6,14,14)
        x=F.max_pool2d(F.relu(self.conv2(x)),2)#(1,6,14,14)->(1,16,10,10)->(1,16,5,5)
        x=x.view(-1,x.size()[1:].numel())#(1,16,5,5)-->(1,16*5*5)
        x=F.relu(self.fc1(x))#(1,16*5*5)->(1,120)
        x=F.relu(self.fc2(x))#(1,120)->(1,84)
        x=self.fc3(x)#(1,84)->(1,10)
        return x

net=Net()
print(net)
a=torch.randn(1,1,32,32)
b=net(a)
print(b.size())

简单的分类程序

环境配置：

conda create -n szh python=3.8
conda install pytorch==1.6.0 cudatoolkit=10.1 -c https