基于PyTorch快速搭建一个简单的CNN卷积神经网络

import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.utils.data as Data
import torchvision
import matplotlib.pyplot as  plt

# Hyper Parameter
EPOCH = 1      # epoch是整个数据集训练次数，即迭代次数
BATCH_SIZE = 50
LR = 0.001
DOWNLOAD_MNIST = False

train_data = torchvision.datasets.MNIST(
    root='./mnist',       # 训练的数据位置
    train=True,
    transform=torchvision.transforms.ToTensor(),     # 将图片像素范围(0-255)--压缩-->(0,1)
    download=DOWNLOAD_MNIST
)

# plot one example （显示指定的训练数据）
# print(train_data.train_data.size())    # 有10000个训练数据(10000, 28, 28)
# print(train_data.train_labels.size())  # 有10000个训练数据，标签(10000)
# plt.imshow(train_data.train_data[3].numpy(),cmap='gray')   # 显示一个指定的训练数据
# plt.title('%i' % train_data.train_labels[0])               # 显示一个指定的训练数据标签
# plt.show()

train_loader = Data.DataLoader(       # 得到一个待批量的生成器
    dataset=train_data,
    batch_size=BATCH_SIZE,
    shuffle=True,
)
# 测试集  因为train=False，所以提取出来的不是train_data(训练数据)，而是test_data(测试数据)
test_data = torchvision.datasets.MNIST(root='./mnist',train=False)
test_x = torch.unsqueeze(test_data.data, dim=1).type(torch.FloatTensor)[:2000]/255. # 将将图片像素除以255使范围(0-255)--压缩-->(0,1)
test_y = test_data.targets[:2000]   # 提取前面2000个数据

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(            # 卷积核(2维)---(过滤器)，作用使提取特征
                in_channels=1,    # 输入图片层数，这里层。(灰度图只有一层，彩色图有RGB三个层)
                out_channels=16,  # 16个不同权重的卷积核，提取16个不同特征（16个厚度），输出16层。不同卷积核*同样输出=不同结果(一般检测16个够用)
                kernel_size=5,    # 卷积核的大小（5*5像素点扫描图片,若长宽不一致(5,3)）
                stride=1,         # 卷积核每隔1个像素(pixel)提取一次特征值
                padding=2,        # 扩展图片像素的圈数 (if stride=1, padding=(kernel_size-1)/2=(5-1)/2=2)
            ),                    # 图片信息 由(1,28,28)--> (16,28,28)，像素不变
            nn.ReLU(),      # 激活函数（还有其他激活函数：nn.Sigmoid, nn.LeakyReLU, nn.Tanh）
            # 池化层参数:池化窗口大小kernel_size，滑动步长stride，padding，卷积对输入数据的空间间隔dilation
            nn.MaxPool2d(kernel_size=2),   # 池化层(最大池化),池化窗口正方形，大小2(筛选信息)。图片信息--> (16,14,14)像素变小一半
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 5, 1, 2),       # 接受上面输出层数16层，然后加工成32层--> (32,14,14)
            nn.ReLU(),                        # 激活函数->(32,14,14)
            nn.MaxPool2d(2)                   # 池化层->(32,7,7),再次缩小一半
        )
        self.out = nn.Linear(32 * 7 * 7, 10)  # 全连接层，把所有特征连接成一条。输出层(32*7*7的输入，10的输出).

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)           # (batch,32,7,7)
        x = x.view(x.size(0), -1)   # (batch,32 * 7 * 7),将全部维度的数据变在一起
        output = self.out(x)
        return output

cnn = CNN()
# print(cnn)    # 打印神经网络的结构(net architecture)

# 优化
optimizer = torch.optim.Adam(cnn.parameters(),lr=LR)  # 用Adam优化器优化cnn网络所有参数(optimizer all cnn parameters)
loss_func = nn.CrossEntropyLoss()       # the target label is not one-hotted

# 开始训练和测试  training and testing
for epoch in range(EPOCH):
    for step, (x, y) in enumerate(train_loader):    # give batch data, normaloze x when it
        output = cnn(x)              # cnn output
        loss = loss_func(output, y)  # cross entropy loss
        optimizer.zero_grad()        # clear gradients for this training step
        loss.backward()              # backpropagation, compute gradients
        optimizer.step()

        # 每50步观察训练效果
        if step % 50 ==0:
            test_output =cnn(test_x)
            pred_y = torch.max(test_output, 1)[1].data.squeeze()
            accuracy = sum(pred_y == test_y) / test_y.size(0)  # 或accuracy = torch.floor_divide(sum(pred_y == test_y), test_y.size(0))
            print('Epoch: ',epoch, '| train loss: %.4f' % loss.item(), '| test accuracy: %.2f' % accuracy)

# print 10 predictions from test data
test_output = cnn(test_x[:10])             # 放前10个数据
pred_y = torch.max(test_output,1)[1].data.numpy().squeeze()   # 得到预测数据
print(pred_y,'prediction number')        # 预测的数值
print(test_y[:10].numpy(),'real number')   # 真实的数值