import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import torchvision
import torchvision.transforms as transforms
import torchvision.datasets as datasets

import time

#归一化处理
## Compose 是创建
transform = transforms.Compose([transforms.ToTensor(),#把数据转换成tensor
                                transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5)) #数据归一化 ，第一组参数是所有channel的平均值,第2组参数是方差
                                                             # 为什么需要平均值的原因是 ，所有图片都有的 ，数据不需要学习，因为平均值都在每张图片里了是干扰项目
                                ])
#训练数据
## 数据在哪
train_data = datasets.CIFAR10('./data',train=True,transform=transform,download=True)
## 数据加载
train_loader = torch.utils.data.DataLoader(dataset=train_data,batch_size=4,shuffle=True,num_workers=2)

#测试数据
## 数据在哪里
test_data =  datasets.CIFAR10('./test',train=False,transform=transform,download=True)
##数据加载
test_loader = torch.utils.data.DataLoader(dataset=test_data,batch_size=4,shuffle=True,num_workers=2)

# min-batch的图片显示
import matplotlib.pyplot as plt
import numpy as np
## 一般显示图片都是numpy array的数据来进行显示的
def imgsshow(img):  ##?  这里有个包升级的问题 ，导致目前没有解决
    img = img/2 + 0.5  #由于在归一化 transform里有去掉平均的处理 在这里需要变回来
    img = img.numpy() # numpay array
    img = np.transpose(img,(1,2,0)) # (c,h,w) -> (h,w,c) 符合正常显示的数据的方式
    plt.show(img)
data_iter = iter(train_loader) #随机加载一个min batch
images , labels = data_iter.next() # 把图片 和 label 分开
# imgsshow(torchvision.utils.make_grid(images))

class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        # LeCun
        # N = (input_size - kernel_size + 2 padding)/stripe +1
        self.conv1 = nn.Sequential(nn.Conv2d(3,6,5,1),  # out =  (32 -5 + 2*0)/1+1 =28
                                   nn.ReLU(),
                                   nn.MaxPool2d(kernel_size=2,stride=2))  # out = 14
        self.conv2 = nn.Sequential(
           nn.Conv2d(6,16,5), # out = (14 -5 +0)1 +1  =10
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2,stride =2)  # out = 10/2  = 5
        )


        self.fc1 = nn.Sequential(
            nn.Linear(16*5*5,120),
            nn.ReLU()
        )
        self.fc2 = nn.Sequential(
            nn.Linear(120,84),
            nn.ReLU()
        )
        self.fc3 = nn.Linear(84,10)
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)

        x = x.view(-1, 16*5*5)
        print ("%%"*50)
        print ('dimision  change',x.shape)
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x
net = Net()
print("current net is ",net)
# 定义损失
criterion = nn.CrossEntropyLoss()
#定义优化器 ，（一旦损失backforward 后，如何更新 weight ，已经更新谁的weight）
tunner = optim.SGD(net.parameters(),lr=0.0001,momentum=0.9)
traning_loss_history = []
test_loss_history = []
for epoch in range(1000):
    net.train()
    running_loss = 0.0
    print ("training.... epoch{0}".format(epoch))
    start_epoc = time.time()
    for i ,data in enumerate(train_loader,0): #一次仅仅提取一个minbatch ，一致到所有的数据取完
        batch_time = time.time()
        imges, labels = data
        imges, labels = Variable(images),Variable(labels)

        tunner.zero_grad()
        outs = net(images)
        loss = criterion(outs, labels)
        # tunner.zero_grad()
        loss.backward()
        tunner.step()
        if i % 1000 == 0:
            print ("i {0} : loss {1} : duration {2}".format(i, loss.item(), (time.time()-batch_time)))
        running_loss += loss.item()
        if i%250 == 0 :
            net.eval()
            with torch.no_grad():
                for data in test_loader:
                    test_images ,test_label = data
                    test_outs = net(test_images)
                    test_loss = criterion(test_outs,test_label)
            traning_loss_history .append(running_loss/250)
            test_loss_history.append(test_loss.item())
            running_loss = 0.0


    print("epoch {0} :: loss {1} :: duration {2}".format(epoch,loss.item(),time.time()-start_epoc))
    # 为什么损失函数会有regression呢 ？ 原因 1 ，minbatch的偶然性导致 ，比如这几次都是飞机 2, learning rrate {}
#画图

plt.figure()
plt.plot(traning_loss_history)
plt.plot(test_loss_history)
plt.legend('training loss','test loss')
plt.tile("Traing /Test loss")
plt.xlabel('#mini_batch *250')
plt.ylabel('Loss')