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
import os
# 批训练加速模型的训练速度
EPOCH = 1 # 定义批次训练的batch数
BATCH_SIZE = 50 # 定义批次训练的batch大小
LR = 0.001 # 学习率，学习率越高模型训练速度越快，但对应得会损失精度
# 下载训练集
train_data = torchvision.datasets.MNIST(
root='./data',
train=True, # 是true下载训练集，如果是false 则下载测试集
transform=torchvision.transforms.ToTensor(), # Converts a PIL.Image or numpy.ndarray to
# torch.FloatTensor of shape (C x H x W) and normalize in the range [0.0, 1.0]
# 载下来的数据转化为tensor格式（0-1），原始数据是pxl（0-255）
download=True,
)
train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)
# 查看数据规模
print(train_data.data.size())
print(train_data.targets.size())
# 查看一张数据
#plt.imshow(train_data.data[0].numpy(), cmap='gray')
#plt.title('%i' % train_data.targets[0])
#plt.show()
# 下载测试集
test_data = torchvision.datasets.MNIST(root='./data', train=False) # 数据还是0-255之间
# print('test_data', test_data.data[:1].shape)
# torch.Size([1, 28, 28])
test_x = torch.unsqueeze(test_data.data, dim=1).type(torch.FloatTensor)[:2000]/255. # 该操作把255数据压缩到0-1之间
# print('test_x', test_x[:1].shape)
# test_x torch.Size([1, 1, 28, 28])
test_y = test_data.targets[:2000]
# 搭建CNN网络模型
# 当我们自己实现类的话，必须继承自nn.Module，并且在init中完成初始化的步骤和forward中完成计算图的前向构建的过程。
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__() # 父类初始化
self.conv1 = nn.Sequential(
# 卷积+激活+池化
# 过滤器 高度-filter 用于提取出卷积出来的特征的属性
# 图片的维度是 (1，28，28) 1 是channel的维度，28x28的宽高
nn.Conv2d(
1, # in_channels=1 输入的通道
16, # out_channels=16 输出的通道(filter的个数)
5, # kernel_size=5，filter的宽高都是5个像素点
1, # stride=1，卷积步长
2, # padding=2填充，如果想输出的和输入的宽高尺寸一样则需要padding=(kernel_size-1)/2
), # (16,28,28)
nn.ReLU(),
# 筛选重要信息 kernel_size=2 把2x2的区域中选最大的值变成1x1
nn.MaxPool2d(kernel_size=2), # (16,14,14)
)
# (16,14,14)
self.conv2 = nn.Sequential(
nn.Conv2d(16, 32, 5, 1, 2), # (32,14,14)
nn.ReLU(),
# 一般用最大值
nn.MaxPool2d(2), # (32,7,7)
)
# 输出层
self.out = nn.Linear(32 * 7 * 7, 10) # （a，b）a是数据维度 b是分类器有十个
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x) # (batch,32,7,7)
x = x.view(x.size(0), -1) # view：将数据变成全连接网络（-1：自动检测矩阵有有多少行，列指定为x.size(0)）
output = self.out(x)
return output
# 模型实例
cnn = CNN()
print(cnn)
'''
# 模型优化和训练
optimizer = torch.optim.Adam(cnn.parameters(), lr=LR) # 优化器
loss_func = nn.CrossEntropyLoss() # 损失函数。项目内容本质上是分类问题，针对多分类问题一般用交叉熵，因为它计算的是一个概率。
# 训练过程
for epoch in range(EPOCH):
for step, (x, y) in enumerate(train_loader):
b_x = Variable(x) # 需要把X，Y使用Variable转化为pytorch能处理的张量形式
b_y = Variable(y)
output = cnn(b_x)
loss = loss_func(output, b_y) # 计算损失函数
optimizer.zero_grad() # 清空过往梯度
loss.backward() # 反向传播，计算当前梯度
optimizer.step() # 根据梯度更新网络参数
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)
print('Epoch: ', epoch, '| train loss: %.4f' % loss.data.numpy(), '| test accuracy: %.2f' % accuracy)
# 保存模型
#torch.save(cnn, './model/cnn.pkl')
'''
# 加载模型
cnn_model = CNN()
cnn_model = torch.load('./model/cnn.pkl')
cnn_model.eval()
# 验证识别成功率
test_output = cnn_model(test_x[:30])
pred_y = torch.max(test_output, 1)[1].data.numpy().squeeze()
print('真实值', test_y[:30].data.numpy().squeeze())
print('预测值', pred_y)
'''
# 验证识别成功率
test_output = cnn(test_x[:10])
# print('test_output 0', test_output)
# print('test_output 1', torch.max(test_output, 1))
# a0 = torch.max(a, dim)
# 其中a为一个tensor
# dim的值为0/1，分别代表索引每列/行最大值
# 返回的值包含两个数据（values, indices) 分别代表最大值的值和所在的索引
# torch.max(a, 1)[1] 代表a中每行最大值的索引
# print('test_output 2', torch.max(test_output, 1)[1].data.numpy())
# numpy.squeeze() 这个函数的作用是去掉矩阵里维度为1的维度。
# 例如 (1, 5)的矩阵经由np.squeeze处理后变成5
# (5, 1, 6)的矩阵经由np.squeeze处理后变成(5, 6)
pred_y = torch.max(test_output, 1)[1].data.numpy().squeeze()
print('真实值', test_y[:10].squeeze())
print('预测值', pred_y)
'''
