#! /usr/bin/env python
# -*- coding:utf-8 -*-
#====#====#====#====
'''
基于Pytorch的LeNet的MNIST手写数字识别Python代码实现
代码步骤
1）导入各种包
2）定义超参数（包括-批大小batch_size、学习速率learing_rate即lr、迭代次数即遍历数据集次数num_epoches等）+ 判断使用cpu还是GPU
3）数据集处理：定义数据集-下载（最好提前下载好！！！），同时建立数据集迭代器-加载数据集（可包括数据集处理）
【为什么用迭代器dataloader】-Pytorch中，训练最好使用迭代器来进行，不然数据集大，内存吃不消
数据集提前下载好后记得download值改成false，另外数据数据集需要组合transform=transforms.Compose([transforms.ToTensor()])
4）定义（使用class定义）并构建网络模型
一部分是网络结构初始化定义：里面一般使用nn.Sequential()顺序网络结构来定义多个卷积+多个全连接模块（激活池化啥的也定义在内）；
一部分是前向传播部分的定义：整个前向传播顺序（多次卷积++平铺+多次全连接+输出）
class定义完成后，记得构建模型-调用下
【为什么使用torch.nn函数】-Pytorch所有的网络在nn包里，用来定义各种模型
【为什么用顺序网络结构Sequential】-简化代码，规范化
【为什么在前向传播部分的卷积核全连接之间使用铺平函数out = out.view(out.size(0), -1)】-将矩阵变成向量传给全连接
5）定义损失函数loss和优化方式SGD /Adm...
6）训练模型（包括-初始化loss和accuracy、for循环训练集dataloader、输入训练集数据、 将输入的tensor类型的训练集img+label转换为variable类型、前向传播-即调用前面构建的模型框架（传入的数据只有训练集img）、计算损失、损失求和、预测最大可能性值对应标签、统计正确率、正确率求和、梯度归零、反向传播、参数更新-通过优化器(计算的是梯度)、打印训练接每个迭代epoch的损失loss和准确率acc）
【传入训练数据集和对应标签集的作用】-训练集是用于前向传播中；标签集用于统计正确率
【为什么将tensor类型的img转换为variable类型】-Pytorch中最重要的就是Variable模块，该模块集成了围绕一个张量所有的操作，包括前向传播、反向传播的各种求偏导数的数值。
7）保存模型
【为什么要保存】-参数已经训练好，达到最优，保存下来后方便训练集拿这个训练好的“数学公式”进行计算并预测值
8）评估测试模型（将模型变换为测试模式-在测试集上测试、for循环的集合是测试集dataloader、初始化loss和accuracy、输入测试数据集、将输入的tensor类型的测试集img+label转换为variable类型、前向传播-即调用前面构建的模型框架（传入的数据只有测试集img）、计算损失、损失求和、预测最大可能性值对应标签、统计正确率、正确率求和、打印正确率和错误率）
【附加】：
1）最好将准确率和误差以曲线的形式打印出来以便调整参数，使用tensorbord包里面的SummaryWriter
代码位置出现在四处:全局变量后面的writer定义+训练模型中两次writer.add-分别代表追加每次迭代时训练集和测试集相关损失和准确率数值+代码末尾writer.close()-释放资源
2）记住：在修改训练集时最重要的是修改数据集名称+class里面几乎每个参数（彩色数据集和黑白数据集的差别在代码上不仅仅在于输入通道数、还有数据集名称+里面各项参数）
3）可以在定义全局变量后直接判断CPU/GPU，以便下面代码方便调用
4) 在网上搜到的代码为了加快运行速度，可以修改迭代次数num_epoches为1/2等，batch_size的不要轻易改动（目前也不是很清楚复杂度，嘻嘻，搞明白了再说吧）
5）训练模型和测试模型的区别在于求导（梯度归0零啦、反向传播啦、参数更新啦这些），是因为训练集需要不断更新参数学习最优模型，测试的时候就可以直接拿来用了
6）可视化输出的正确率和错误率曲线会发现有震荡现象，是因为收敛不到最优值，所以有震荡，但正确率整体上升、错误率整体下降（会发现最开始阶段变化最明显）
7）使用变量时名字一定要避开关键词，如定义全局变量迭代次数时，可使用num_epoches等
8）函数调用时以防自己记错参数位置，最好直接使用“键=值”形式！！！
9）在定义类时关于里面参数如何选择，我也不是很清楚？？？，但可以肯定得是卷积核个数-即每次输出通道都是自己人为设定的
10）关于注释-单行#，多行则使用''' '''
11）在此强调一下在你的那个代码里面，你的测试模型部分放在了定义的def Accuracy()函数里面哈！！！
【代码详细解释参考】
https://zhuanlan.zhihu.com/p/30117574 重要
https://blog.csdn.net/xiaoheizi_du/article/details/88571240 重要
https://blog.csdn.net/l2181265/article/details/95351618
【其他可学习】
CIFAR10+LeNet
https://blog.csdn.net/l2181265/article/details/95351618
'''
#====#====#====#====
###### 第一步、导入代码所需库函数 ######
import torch
from torch import nn, optim # nn是神经网络模块
#import torch.nn.functional as F # 有的代码将torch.nn简化为F使用
from torch.autograd import Variable # Variable是Pytorch数据格式模块
from torch.utils.data import DataLoader # 将torch.utils.data简化，所以下面加载数据时可直接DataLoader.XXX
from torchvision import transforms # torchvision是Pytorch的外围库，该库包含了各种关于图像的各种功能函数
from torchvision import datasets # 将torchvision简化，所以下面下载和处理数据时可直接使用dtasets.XXX
#from logger import Logger #其中logger包是记录日志用的，并不需要（本代码进行了编写只不过在最后一步的注释部分）
###### 第二步、定义超参数 #####
'''
我们使用的batch_size为64，有的人好奇为什么使用64，或者32，我的理解是这样的，当我们的数据大小为2的幂次数，计算机会计算的特别快，
因为计算机是二进制嘛，如果是2的幂次数的话，计算机很容易通过移位来进行计算。
'''
batch_size = 128 # 批的大小（指将数据集分成n个batch(批/块)，每一个batch的大小）
learning_rate = 1e-3 # 学习速率（一般以缩写lr表示；1e-3=0.001）
num_epoches = 1 # 遍历整个训练集的次数，此处为了节约时间，训练一次（因为epoch是一个关键字，所以这里表示成num_epoches）
# 数据类型转换，转换成numpy类型
#def to_np(x):
# return x.cpu().data.numpy()
##### 第三步、定义训练+测试数据集并下载训练+测试集 MNIST 手写数字训练集 #####
'''
数据集下载和处理使用说明
Pytorch的torchvision.datasets包含MNIST、COCO、LSUN、ImageFolder、Imagenet-12、CIFAR、STL10等数据集。
可以通过datasets.MNIST这种方式来下载调用这些数据集
函数使用-datasets.MNIST(root, train, transform, target_transform, download)参数说明：
root： 数据集的目录名，即数据集下载和保存的地址
train：train=True表示训练集，train=False表示测试集
transform：需要转换成张量形式，即transform=transforms.ToTensor()
target_transform：一般不用
download：download=True表示从网上下载数据集，download=False表示数据已经下载过。
'''
train_dataset = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor())
##### 第三步、建立数据集迭代器---定义训练批处理数据并加载数据集 #####
'''
数据加载器函数说明---由于第一步使用了from torch.utils.data import DataLoader 因此简化了该写法--直接使用DataLoader替代torch.utils.data.DataLoader
torch.utils.data.DataLoader(dataset, batch_size, shuffle, sampler, num_workers, collate_fn, pin_memory, drop_last)
一般常用dataset+batch_size+shuffle+num_workers这四个！！！！
dataset (Dataset)：加载数据的数据集
batch_size (int, optional)：数据加载，将每X个拼成一个批次，即每个batch加载多少个样本(默认: 1)
shuffle (bool, optional)：设置为True时会在每个epoch重新打乱数据(默认: False).shuffle只是为了保证随机性而已，这只在训练的时候有作用！！！！。
sampler (Sampler, optional)：不常用 -定义从数据集中提取样本的策略。如果指定，则忽略shuffle参数。
num_workers (int, optional)：加载数据时使用多少子进程。默认值为0，表示在主进程中加载数据。
collate_fn (callable, optional)：不常用
pin_memory (bool, optional)：不常用
drop_last (bool, optional)：不常用-如果数据集大小不能被batch size整除，则设置为True后可删除最后一个不完整的batch。如果设为False并且数据集的大小不能被batch size整除，则最后一个batch将更小。(默认: False)。
'''
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)# shuffle只是为了保证随机性而已，这只在训练的时候有作用
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
##### 第四步 定义LeNet卷积神经网络模型 #####
'''
框架（2次卷积+3次全连接/3次卷积+2次全连接）
一般以类Class来定义，里面包含模型初始化部分（卷积模块+全连接模块）和前向传播2大部分
模型初始化部分def __init__(self):
super的用法：LeNet继承父类nn.Model的属性，并用父类的方法初始化这些属性
nn.Sequential()：顺序网络结构、即一个时序容器,Modules 会以他们传入的顺序被添加到容器中。这个容器里可以初始化卷积层、激活层和池化层。
卷积函数：nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)：用来初始化卷积层（包括激活、池化等一些操作）
in_channels：输入图像的通道深度--灰色图片是1，彩色是3
out_channels：输出图像的通道深度--卷积核个数
kernel_size：卷积核的尺寸大小
stride：步长--每隔多少像素跳一下，默认为1
padding：边缘填充--在图像周围一圈补0，为了扫描完整（有时候不整除）,默认为0
dilation：膨胀卷积，默认为1，目前用不到
groups：组卷积，默认为1 ，目前用不到
bias：默认为True
这些参数顺序不用刻意记住，一般也就是用这样的顺序--输入通道+输出通道+步长+填充（假如没记住，可以使用键=值的形式）！！！！
如nn.Conv2d(in_channels = 1,out_channels = 6,kernel_size = 5,stride=1,padding=0)
池化函数： nn.MaxPool2d(kernel_size, stride)
kernel_size：卷积核的尺寸大小
stride：步长--每隔多少像素跳一下
因为一般只用到这两个参数，因此可直接传值的形式而不用键=值形式！！！
全连接函数：nn.Linear(in_features, out_features, bias=True)
in_features：每个输入样本的大小
out_features：每个输出样本的大小
bias：用不到，若设置为False，这层不会学习偏置。默认值：True
前向传播部分def forward(self, x)
直接调用多次卷积
铺平
调用多次全连接
有时这里会有个softmax层调用（要根据后面的损失函数来决定，损失函数是交叉熵损失，那么这里就不调用了）！！！我也不是很懂为啥，可以暂时记住就这么做
返回输出
'''
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv1 = nn.Sequential( # input_size=(1*28*28)
# 需要注意的是由于LeNet处理的默认输入时32*32的图片，这里加padding=2，即上下左右各padding 2个单位像素，扩充到32*32
# 【卷积核个数-即每次输出通道都是自己人为设定的】
nn.Conv2d(1, 6, 5, 1, 2), # output_size=(6*28*28) 此处padding=2保证输入输出尺寸相同--(28-5+2*2)/1+1=28
nn.ReLU(), # 先激活 # output_size=(6*28*28) 保持维度不变
nn.MaxPool2d(2, 2) # 在 2x2 空间里向下采样取最有效信息, output shape (6, 14, 14)--(28-2+2*0)/2+1=14
)
self.conv2 = nn.Sequential( # input_size=(6*14*14)
nn.Conv2d(6, 16, 5), # output_size=(16*10*10) --(14-5+2*0)/1+1=10
nn.ReLU(), # 先激活 output_size=(16*10*10) 保持维度不变
nn.MaxPool2d(2, 2) # 在 2x2 空间里向下采样取最有效信息, output shape (16, 5, 5)--（10-2+2*0)/2+1=5
)
self.fc1 = nn.Sequential(
nn.Linear(16 * 5 * 5, 120), # 120是本层卷积核个数，此处第一个参数必须为16*5*5
nn.BatchNorm1d(120), # 加快收敛速度的方法（注：批标准化一般放在全连接层后面，激活函数层的前面）
nn.ReLU()
)
'''
【为啥设置84】输出层由欧式径向基函数（Euclidean Radial Basis Function）单元组成，每类一个单元，每个有84个输入。换句话说，
每个输出RBF单元计算输入向量和参数向量之间的欧式距离。输入离参数向量越远，RBF输出的越大。
一个RBF输出可以被理解为衡量输入模式和与RBF相关联类的一个模型的匹配程度的惩罚项。
用概率术语来说，RBF输出可以被理解为F6层配置空间的高斯分布的负log-likelihood。
给定一个输入模式，损失函数应能使得F6的配置与RBF参数向量（即模式的期望分类）足够接近。
这些单元的参数是人工选取并保持固定的（至少初始时候如此）。这些参数向量的成分被设为-1或1。虽然这些参数可以以-1和1等概率的方式任选，或者构成一个纠错码，
但是被设计成一个相应字符类的7*12大小（即84）的格式化图片。这种表示对识别单独的数字不是很有用，但是对识别可打印ASCII集中的字符串很有用。
'''
self.fc2 = nn.Sequential(
nn.Linear(120, 84),
nn.BatchNorm1d(84),
nn.ReLU()
)
self.fc3 = nn.Linear(84, 10)
# self.sfx = nn.Softmax() #因为后面使用的交叉损失熵，所以没有softmax层，我目前也不知道为啥？？？？
def forward(self, x):
out = self.conv1(x)
out = self.conv2(out)
# print(x.shape)
out = out.view(out.size(0), -1)# nn.Linear()的输入输出都是维度为一的值，所以要把多维度的tensor展平成一维
out = self.fc1(out)
out = self.fc2(out)
out = self.fc3(out)
# out = self.fc1(out)# .Net最后一层为什么没有softmax？因为后面用的loss函数是CrossEntropyLoss，不需要softmax
return out
##### 第四步、构建LeNet模型 #####
model = LeNet() # 图片大小是28x28,输入深度是1，最终输出的10类
##### 判断是否有GPU加速 ####
use_gpu = torch.cuda.is_available()
if use_gpu:
model = model.cuda()
print('USE GPU')
else:
print('USE CPU')
##### 第五步、定义损失函数loss（此处使用的是交叉损失熵）和优化器（此处使用的是SGD） ######
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)
#logger = Logger('./logs')
##### 第六步、训练模型 #####
for epoch in range(num_epoches): # 设定整个数据集的迭代次数
print('epoch {}'.format(epoch + 1)) # .format为输出格式，formet括号里的即为左边花括号的输出
print('*' * 10)
running_loss = 0.0 # 初始化训练集的损失值loss
running_acc = 0.0 # 初始化训练集的准确率值accuracy
for i, data in enumerate(train_loader, 1): # data是一个以两个张量为元素的列表：两个元素分别是img和label，元素类型为tensor类型。
##### 输入数据并处理 ##### # enumerate枚举训练器中数据，参考[enumerate](https://www.runoob.com/python/python-func-enumerate.html)
img, label = data # 输入训练集数据（因为for循环的是训练集，因此这里就是训练集的img和label)
# cuda
if use_gpu: # 使用GPU的话还需要将数据集传到GPU
img = img.cuda()
label = label.cuda()
img = Variable(img) # 将tensor类型的img转换为variable类型。
label = Variable(label) # 将tensor类型的label转换为variable类型。
##### 向前传播 #####
out = model(img) # 前向传播-调用构建的模型，使用的是img数据集
loss = criterion(out, label) # 计算交叉熵损失
running_loss += loss.item() * label.size(0) # 损失求和
_, pred = torch.max(out, 1) # 预测最大值所在的位置标签，即预测数字是0-9的几
'''
torch.max()：返回输入张量所有元素的最大值。
torch.max(input, dim, max=None, max_indices=None)：返回输入张量给定维度上每行的最大值，并同时返回每个最大值的位置索引。
输出形状中，将dim维设定为1，其它与输入形状保持一致。
input (Tensor)：输入张量
dim (int)：指定的维度
max (Tensor, optional)：结果张量，包含给定维度上的最大值
max_indices (LongTensor, optional)：结果张量，包含给定维度上每个最大值的位置索引,预测最大值所在的位置标签，即预测的数字。
'''
num_correct = (pred == label).sum() # 统计预测正确的数目
accuracy = (pred == label).float().mean() # 计算准确率（平均值）
running_acc += num_correct.item() # 准确率求和
##### 后向传播 #####
optimizer.zero_grad() # 梯度（即导数）归零---zero_grad()：清空所有被优化过的Variable的梯度.
loss.backward() # 反向传播
optimizer.step() # 参数更新-通过梯度做一步参数更新---step()：进行单次优化
"""
# ========================= Log ======================
step = epoch * len(train_loader) + i
# (1) Log the scalar values
info = {'loss': loss.data[0], 'accuracy': accuracy.data[0]}
for tag, value in info.items():
logger.scalar_summary(tag, value, step)
# (2) Log values and gradients of the parameters (histogram)
for tag, value in model.named_parameters():
tag = tag.replace('.', '/')
logger.histo_summary(tag, to_np(value), step)
logger.histo_summary(tag + '/grad', to_np(value.grad), step)
# (3) Log the images
info = {'images': to_np(img.view(-1, 28, 28)[:10])}
for tag, images in info.items():
logger.image_summary(tag, images, step)
if i % 300 == 0:
print('[{}/{}] Loss: {:.6f}, Acc: {:.6f}'.format(
epoch + 1, num_epoches, running_loss / (batch_size * i),
running_acc / (batch_size * i)))
"""
##### 打印每个epoch的loss和acc ######
print('Finish {} epoch, Loss: {:.6f}, Acc: {:.6f}'.format(
epoch + 1, running_loss / (len(train_dataset)), running_acc / (len(train_dataset)))) # 这里又有.format()的用法。
##### 测试（评估）模型 #####
model.eval() # 开始验证测试集
eval_loss = 0 # 初始化测试集错误率
eval_acc = 0 # 初始化测试准确率
for data in test_loader: # 测试的是测试数据集
img, label = data # 输入测试集数据（因为for循环的是测试集，因此这里就是测试集的img和label)
if use_gpu:
img = Variable(img, volatile=True).cuda()
label = Variable(label, volatile=True).cuda()
else:
img = Variable(img, volatile=True)
label = Variable(label, volatile=True)
out = model(img) # 前向传播-调用构建的模型，使用的是img数据集
# print("output_test:",out.shape)
loss = criterion(out, label) # 计算损失
eval_loss += loss.item() * label.size(0) # 损失求和
_, pred = torch.max(out, 1) # 预测最大值所在的位置标签，即预测数字是0-9的几
# print("predicted:",pred.shape)
num_correct = (pred == label).sum() # 统计预测准确个数
eval_acc += num_correct.item() # 准确率求和
print('Test Loss: {:.6f}, Acc: {:.6f}'.format(eval_loss / (len(
test_dataset)), eval_acc / (len(test_dataset)))) # 打印
print()
##### 保存模型 ######
'''
torch.save(model.state_dict(), './cnn.pth')
model.state_dict()：需要保存的模型
'./cnn.pth'：模型的名称
'''
torch.save(model.state_dict(), './LeNet.pth')
