import  torch
import torchvision#包含一些数据集，像如mnist
import torch.utils.data as Data
import torch.nn as nn
import matplotlib.pyplot as plt

#超参数
EPOCH=1
BATCH_SIZE=50
LR=0.01
DOWNLOAD_MNIST=True#是否下载数据集


#mnist手写数字数据集
train_data=torchvision.datasets.MNIST(
    root="mnist",
    train=True,#是否是训练
    transform=torchvision.transforms.ToTensor(),#将图像或numpy数组数据形式转化成tensor的floattensor(C*H*M)形式
    download=DOWNLOAD_MNIST #初始下载数据集用，后面就不用了
)

#测试集
test_data=torchvision.datasets.MNIST("mnist",train=False)
#批训练和dataloader数据封装
train_loader=Data.DataLoader(
    dataset=train_data,
    batch_size=BATCH_SIZE,
    shuffle=True
    
)

#测试数据  只取前2000个
test_x = torch.unsqueeze(test_data.test_data, dim=1).type(torch.FloatTensor)[:2000]/255.   # shape from (2000, 28, 28) to (2000, 1, 28, 28), value in range(0,1)
test_y = test_data.test_labels[:2000]

#实现卷积层模型
class CNN(torch.nn.Module):
    def __init__(self):
        super(CNN,self).__init__()
        #第一层卷积
        self.conv1=torch.nn.Sequential(
            torch.nn.Conv2d(
                in_channels=1,
                out_channels=16,
                kernel_size=5,
                stride=1,
                padding=2#padding=(kernel_size-1)/2时，卷积出来的图片尺寸不变
            ),#此时输入是(,28*28)输出是(16*28*28)
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2)#此时图片尺寸减小一半  16*28*28--》16*14*14
        )
        
        #第二层卷积
        self.conv2=torch.nn.Sequential(
            torch.nn.Conv2d(16,32,5,1,2),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2)
        )#输入是16*14*14  输出是32*7*7
        
        #添加全连接层
        self.output=torch.nn.Linear(32*7*7,10)
        
        
    #前向传播  卷积
    def forward(self,x):
        x=self.conv1(x)
        x=self.conv2(x)
        #将此时的32*7*7的图片铺平  拉成一维32*7*7的数组形式
        x=x.view(x.size(0),-1)# 展平多维的卷积图成 (batch_size, 32 * 7 * 7)
        output=self.output(x)
        return output
        
        
        

#实例化卷积神经网络
cnn=CNN()
print(cnn)#打印网络结构

#训练

#优化器
optimizer=torch.optim.Adam(cnn.parameters(),lr=LR)
#损失函数
loss_func=torch.nn.CrossEntropyLoss()#用交叉熵损失函数  标签是one-hot编码

#训练和测试
for epoch in range(EPOCH):
    #批训练
    for step,(batch_x,batch_y) in enumerate(train_loader):
        output=cnn(batch_x)
        loss=loss_func(output,batch_y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()


#预测
test_output = cnn(test_x[:10])
pred_y = torch.max(test_output,1)[1].data.numpy().squeeze()
print(pred_y, 'prediction number')
print(test_y[:10].numpy(), 'real number')