【Pytorch】LeNet的pytorch写法

最近在学习pytorch的内容，写几篇博客记录一下学习过程。

熟悉pytorch的朋友知道nn.Module是nn中最重要的类，可以把它看作一个网络的封装，包含网络各层定义及forward方法，调用forward(input)方法，可返回前向传播的结果。

从最早的卷积神经网络LeNet为例，看看如何用nn.Module实现。LeNet的网络结构如图所示。

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.autograd import  Variable

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 1代表输入图片未单通道，6表示输出通道数，5表示卷积核为5*5
        self.conv1 = nn.Conv2d(1, 6, 5) # 这里只能写三个参数，但是要理解输入其实是4维的，也必须是4维的
        # 卷积层
        self.conv2 = nn.Conv2d(6, 16, 5)
        # 仿射层/全连接层,y=Wx+b
        self.fc1 = nn.Linear(16*5*5, 120)
        print(666)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
        # self.fc4 = nn.Linear(10,1) #zhouz


    def forward(self, x):
        # 卷积->激活->池化
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        # reshape '-1'表示自适应
        x = x.view(x.size()[0], -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        # x = F.relu(self.fc3(x)) #zhouz 为了与原书中代码保持一致，暂时改回
        x = self.fc3(x)
        return x
        # print(7777)


net = Net().cuda()
# print(net)

# '测试输入为tensor时，输出size为(1,10)'
# input = Variable(torch.arange(0,10))
input = Variable(torch.randn(1, 1, 32, 32)).cuda()
# out = net(input)
# print(out)
# print(out.size())

# '第一次测试结束'

# '第二次测试开始'

# for name,parameters in net.named_parameters():
#     print(name,':',parameters.size())

# '第二次测试结束'

# '第三次测试开始' 输出网络的可学习参数

params = list(net.parameters())
print(len(params))

# '第三次测试结束'

# '第四次测试开始'
# net.zero_grad()
# out.backward(Variable(torch.ones(1,10))) # 反向传播
# '第四次测试结束'

output = net(input)
target = Variable(torch.arange(1, 11)).float().cuda()
print(output.type == target.type)
print('output', output)
# print(target)

criterion = nn.MSELoss()
loss = criterion(output, target)
print(loss)

net.zero_grad()
print('反向传播之前的conv1.bias的梯度:')
print(net.conv1.bias.grad)
loss.backward()

print('反向传播之后conv1.bias的梯度:')
print(net.conv1.bias.grad)

# 在cpu跑会出错，在gpu跑能运行

 

调试的时候遇到了arange和range的问题，参考以下链接：https://www.jianshu.com/p/438c09259220

【12.27补充】 我在cpu上测试会出问题，拿到ubuntu上跑代码能正常运行，要改的地方包括float()和cuda()，改了之后就能直接运行了。