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') # 真实的数值
