pytorch cnn卷积神经网络代码详解 注释详细

"""---------------------cnn卷积神经网络---------------------"""
import torch
from torch import nn, optim
import torch.nn.functional as F
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision import datasets
from log1 import Logged

# 定义超参数
batch_size = 128
learning_rate = 1e-2
num_epoches = 20

"""04-Convolutional Neural Network"""
def to_np(x):
    return x.cpu().data.numpy()


# 下载训练集 MNIST 手写数字训练集
train_dataset = datasets.MNIST(
    root='F:/PycharmProjects/pytorch-beginner-master/02-Logistic Regression/data', train=False, transform=transforms.ToTensor(), download=True)

test_dataset = datasets.MNIST(
    root='F:/PycharmProjects/pytorch-beginner-master/02-Logistic Regression/data', train=False, transform=transforms.ToTensor())

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)


#定义 Convolution Network 模型
#nn.Sequential：这个表示将一个有序的模块写在一起，也就相当于将神经网络的层按顺序放在一起，这样可以方便结构显示
class Cnn(nn.Module):
    def __init__(self, in_dim, n_class):
        super(Cnn, self).__init__()
        self.conv = nn.Sequential(
            # 开始输入[batch_size=128, in_dim=1, 28,28]经过第一层Conv2d后，用公式W2 = ((input-kernel_size+2padding)/stride) + 1计算输出为：[128,6,28,28]
            nn.Conv2d(in_channels=in_dim, out_channels=6, kernel_size=3, stride=1, padding=1),
            # 修正线性单元，是神经元的激活函数 默认设置为False，表示新创建一个对象对其修改，也可以设置为True，表示直接对这个对象进行修改
            nn.ReLU(True),#输入[128,6,28,28],输出为：[128,6,28,28]
            # 池化层可以非常有效地缩小矩阵的尺寸。从而减少最后全连接层中的参数
            nn.MaxPool2d(2, 2),#输入：[128,6,28,28]输出：[128,6,14,14]
            nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0),#输入：[128,6,14,14]输出：[128,16,10,10]
            nn.ReLU(True), #
        nn.MaxPool2d(kernel_size=2, stride=2))#输入：[128,16,10,10]输出：[128,16,5,5]

    self.fc = nn.Sequential(
        nn.Linear(400, 120), #输入为：[128,400]，输出为：[128,120]
        nn.Linear(120, 84), #输入为：[128,120]，输出为：[128,84]
        nn.Linear(84, n_class))#输入为：[128,84]，输出为：[128,10]

def forward(self, x):
    out = self.conv(x)#输入x=[128,1,28,28],上面__init__一步步说明了out=[128,16,5,5]的由来
    #返回一个有相同数据但大小不同的tensor,-1就代表这个位置由其他位置的数字来推断
    out = out.view(out.size(0), -1)#out变为[128,x] x=128*16*5*5/128，所以x值为400，out为[128,400]
    out = self.fc(out)#进入这个函数看,输出为[128,10]
    return out


model = Cnn(in_dim=1, n_class=10)  # 图片大小是28x28
use_gpu = torch.cuda.is_available()  # 判断是否有GPU加速
if use_gpu:
    model = model.cuda()
# 定义loss和optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)
#logger = Logged('F:/PycharmProjects/pytorch-beginner-master/04-Convolutional Neural Network/logs')
#开始训练
for epoch in range(num_epoches):
    print('epoch {}'.format(epoch + 1))
    print('*' * 10)
    for i, data in enumerate(train_loader, 1):
        img, label = data
        if use_gpu:
            img = img.cuda()
            label = label.cuda()
        img = Variable(img)
        label = Variable(label)
        # 向前传播

    # 开始输入为[batch_size=128, in_dim=1, 28,28] 其分别代表：batch_size组数据，通道数为1，
    # 高度为28，宽为28,输出为[batch_size=128,n_class=10],详情进model函数看
    out = model(img)
    loss = criterion(out, label)

    # 向后传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if i % 20 == 0:
        print("loss: ", loss.data.item())


model.eval()
for data in test_loader:
    img, label = data
    if use_gpu:
        with torch.no_grad():
            img = Variable(img).cuda()
            label = Variable(label).cuda()
    else:
        with torch.no_grad():
            img = Variable(img)
            label = Variable(label)
    out = model(img)
    loss = criterion(out, label)
print("test loss: ", loss.data.item())

# 保存模型
torch.save(model.state_dict(), 'F:/PycharmProjects/pytorch-beginner-master/04-Convolutional Neural Network/cnn.pth')