# -*- coding: utf-8 -*-
"""
Created on Mon Sep  3 08:38:27 2018

@author: www
"""

import torch
from torch import nn
from torchvision.datasets import MNIST
import numpy as np
from torch.autograd import Variable

train_set = MNIST('./data', train=True, download=True)
test_set = MNIST('./data', train=False, download=True)

#觀察一下資料
a_data, a_label = train_set[0]
print(a_data)
print(a_label)

#這裡的讀入的資料是 PIL 庫中的格式，我們可以非常方便地將其轉換為 numpy array
a_data = np.array(a_data, dtype='float32')
print(a_data)

#對於神經網路，我們第一層的輸入就是 28 x 28 = 784，所以必須將得到的資料
#我們做一個變換，使用 reshape 將他們拉平成一個一維向量
def get_data(x):
     x = np.array(x, dtype='float32')/255
     x = (x-0.5)/0.5
     x = x.reshape((-1,))
     x = torch.from_numpy(x)
     return x
train_set = MNIST('./data', train=True, transform=get_data, download=True)
test_set = MNIST('./data', train=False, transform=get_data, download=True)

a, a_label = train_set[0]
print(a.shape)
print(a_label)

#使用pytorch自帶的DataLoader定義一個數據迭代器
#使用這樣的資料迭代器是非常有必要的，如果資料量太大，就無法一次將他們全部讀入記憶體，
#所以需要使用 python 迭代器，每次生成一個批次的資料
from torch.utils.data import DataLoader
train_data = DataLoader(train_set, batch_size=64, shuffle=True)
test_data = DataLoader(test_set, batch_size=128, shuffle=False)

a, a_label = next(iter(train_data))

#打印出一個批次的資料
print(a.shape)
print(a_label.shape)

#使用Sequential定義4層神經網路
net = nn.Sequential(
     nn.Linear(784, 400),
     nn.ReLU(),
     nn.Linear(400, 200),
     nn.ReLU(),
     nn.Linear(200, 100),
     nn.ReLU(),
     nn.Linear(100, 10)
)

#交叉熵在 pytorch 中已經內建了，交叉熵的數值穩定性更差，所以內建的函式已經幫我們解決了這個問題
#定義loss函式
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(net.parameters(), 1e-1)

#開始訓練
losses = []
acces = []
eval_losses = []
eval_acces = []

for e in range(20):
     train_loss = 0
     train_acc = 0
     net.train()
     for im, label in train_data:
          im = Variable(im)
          label = Variable(label)
          #前向傳播
          out = net(im)
          loss = criterion(out, label)
          #反向傳播
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
          #記錄誤差
          train_loss += loss.item()
          #計算分類的準確率
          _, pred = out.max(1)
          num_correct = (pred==label).sum().item()
          acc = num_correct / im.shape[0]
          train_acc += acc
     losses.append(train_loss / len(train_data))
     acces.append(train_acc / len(train_data))
     #在測試集上檢驗效果
     eval_loss = 0
     eval_acc = 0
     net.eval() 
     for im, label in test_data:
          im = Variable(im)
          label = Variable(label)
          out = net(im)
          loss = criterion(out, label)
          #記錄誤差
          eval_loss += loss.item()
          #記錄準確率
          _, pred = out. max(1)
          num_correct = (pred==label).sum().item()
          acc = num_correct / im.shape[0]
          eval_acc += acc
     eval_losses.append(eval_loss / len(test_data))
     eval_acces.append(eval_acc / len(test_data))
     print('epoch:{}, Train Loss:{:.6f}, Train Acc:{:.6f}, Eval Loss:{:.6f}, Eval Acc:{:.6f}'
     .format(e, train_loss/len(train_data), train_acc/len(train_data), 
             eval_loss/len(test_data), eval_acc/len(test_data)))

#最後一次的資料
#epoch:19, Train Loss:0.008779, Train Acc:0.997385, Eval Loss:0.072433, Eval Acc:0.982793

#畫出loss曲線和準確率曲線
import matplotlib.pyplot as plt
plt.title('train loss')
plt.plot(np.arange(len(losses)), losses)

plt.plot(np.arange(len(acces)), acces)
plt.title('train acc')

plt.plot(np.arange(len(eval_losses)), eval_losses)
plt.title('test loss')

plt.plot(np.arange(len(eval_acces)), eval_acces)
plt.title('test acc')