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


#模型训练超参数设置，构建训练模型
EPOCH = 2 #两遍迭代
BATCH_SIZE = 50
LR = 0.001 #学习速度
DOWNLOAD_MNIST = True #如果没有源数据，则DOWNLOAD_MNIST = True

train_data = torchvision.datasets.MNIST(
    root='./mnist',
    train=True,
    download=DOWNLOAD_MNIST
)

#数据下载后是不可以直接看的，查看第一张图片数据
print(train_data.data.size())  #显示数量，以及维度
print(train_data.targets.size()) #显示数量
print(train_data.data[0]) #查看第一张


# 画一个图片显示出来
plt.imshow(train_data.data[0].numpy(),cmap='gray')
plt.title('%i'%train_data.targets[0])
plt.show()

# 训练和测试数据准备
train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)
#DataLoader会根据参数生成batch数据
test_data = torchvision.datasets.MNIST(
    root='./mnist',
    train=False,
)

# 这里只取前3千个数据吧，差不多已经够用了，然后将其归一化。
with torch.no_grad(): #作用：在该模块下，所有计算出的tensor的requires_grad都自动设置为false
    test_x = Variable(torch.unsqueeze(test_data.data, dim=1)).type(torch.FloatTensor)[:3000] / 255
    test_y = test_data.targets[:3000]
# torch.unsqueeze() 函数起到升维的作用，dim等于几表示在第几维度加1，比如x.size为[2,3],x.unsqueeze(0)之后，size变为[2,1,3]
# variable就是变量的意思，它是一种可以变化的常量，符合人类反向传播，参数更新的特性
'''开始建立CNN网络''',


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__() #需要调用父类的构造方法；
        '''
        一般来说，卷积网络包括以下内容：
        1.卷积层
        2.神经网络
        3.池化层
        '''
        # nn.sequential 一个有序的容器，神经网络模块将按照在传入构造器的顺序一次被添加到计算图中执行。
        self.conv1 = nn.Sequential(
            #nn.conv2d 二维卷积，先实例化再使用。在pytorch的nn模型中，它不需要手动定义网络层的权重和偏置。
            nn.Conv2d(  # 输入 (1,28,28)
                in_channels=1,  # 传入的图片是几层的，灰色为1层，RGB为三层
                out_channels=16,  # 输出的图片是几层
                kernel_size=5,  # 代表扫描的区域点为5*5，表示卷积核的大小；若为3*5的卷积核，就写为kernel_size = (3,5)
                stride=1,  # filter step卷积核在图像窗口上每次平移的间隔，即所谓的步长
                padding=2,  # 边框补全，其计算公式=（kernel_size-1）/2=(5-1)/2=2
            ),  # 2d代表二维卷积           输出 (16,28,28)
            nn.ReLU(),  # 非线性激活层，分段性函数，把所有的负值都变为0，而正值不变，即单侧抑制
            nn.MaxPool2d(kernel_size=2),  # 设定这里的扫描区域为2*2，且取出该2*2中的最大值 2*2采样，28/2=14，输出为 (16,14,14)
        )

        self.conv2 = nn.Sequential(
            nn.Conv2d(  # 输入 (16,14,14)
                in_channels=16,  # 这里的输入是上层的输出为16层
                out_channels=32,  # 在这里我们需要将其输出为32层
                kernel_size=5,  # 代表扫描的区域点为5*5
                stride=1,  # 就是每隔多少步跳一下
                padding=2,  # 边框补全，其计算公式=（kernel_size-1）/2=(5-1)/2=
            ),  # 输出(32,14,14)
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),  # 设定这里的扫描区域为2*2，且取出该2*2中的最大值 输出为(32,7,7)，这里是三维数据
        )

        self.out = nn.Linear(32 * 7 * 7, 10)  # 全连接层7*7*32 注意一下这里的数据是二维的数据

# 需要·实现·forward()方法，用于网络的前向传播，而反向传播只需要·调用·Variable.backward()即可。

    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)
        #view()函数的功能和reshape类似，用来转换size的大小。
        output = self.out(x) #全连接层分类器
        return output

# 查看网络结构
cnn = CNN()
print(cnn) # 使用print(cnn)可以看到网络的结构详细信息，可以看到ReLU()也是一层layer


# 添加优化方法
optimizer = torch.optim.Adam(cnn.parameters(), lr=LR) #Pytorch中的一种优化器
# 指定损失函数使用交叉信息熵
loss_fn = nn.CrossEntropyLoss()
#信息量：它是用来衡量一个时间的不确定性的；一个事件发生的概率越大，不确定性越小，则它所携带的信息量越小
#熵：它用来衡量一个系统的混乱程度的，代表一个系统中信息量的总和；信息总和越大，表明这个系统不确定性就越大。
#交叉熵：它主要刻画的是实际输出（概率）与期望输出（概率）的距离，交叉熵的值越小，两个概率分布越近。

'''
开始训练我们的模型哦
'''
step = 0
for epoch in range(EPOCH):
    # 加载训练数据
    for step, data in enumerate(train_loader):
        x, y = data
        # 分别得到训练数据的x和y的取值
        b_x = Variable(x)
        b_y = Variable(y)

        output = cnn(b_x)  # 调用模型预测
        loss = loss_fn(output, b_y)  # 计算损失值
        optimizer.zero_grad()  # 每一次循环之前，将梯度清零
        loss.backward()  # 反向传播
        optimizer.step()  # 梯度下降

        # 每执行50次，输出一下当前epoch、loss、accuracy
        if (step % 50 == 0):
            # 计算一下模型预测正确率
            test_output = cnn(test_x)
            y_pred = torch.max(test_output, 1)[1].data.squeeze()
            accuracy = sum(y_pred == test_y).item() / test_y.size(0)

            print('now epoch :  ', epoch, '   |  loss : %.4f ' % loss.item(), '     |   accuracy :   ', accuracy)

'''
打印十个测试集的结果
'''
test_output = cnn(test_x[:10])
y_pred = torch.max(test_output, 1)[1].data.squeeze()  # 选取最大可能的数值所在的位置
print(y_pred.tolist(), 'predecton Result')
print(test_y[:10].tolist(), 'Real Result')