#需要導(dǎo)入的包
import torch
import torch.nn as nn#用于構(gòu)建網(wǎng)絡(luò)層
import torch.optim as optim#導(dǎo)入優(yōu)化器
from torch.utils.data import DataLoader#加載數(shù)據(jù)集的迭代器
from torchvision import datasets, transforms#用于加載mnsit數(shù)據(jù)集

#下載數(shù)據(jù)集

train_set = datasets.MNIST('./data', train=True, download=True,transform = transforms.Compose([
         transforms.ToTensor(),
         transforms.Normalize((0.1037,), (0.3081,))
       ]))
test_set = datasets.MNIST('./data', train=False, download=True,transform = transforms.Compose([
         transforms.ToTensor(),
         transforms.Normalize((0.1037,), (0.3081,))
       ]))

#構(gòu)建網(wǎng)絡(luò)（網(wǎng)絡(luò)結(jié)構(gòu)對應(yīng)tensorflow的那一篇文章）

class Net(nn.Module):

  def __init__(self, num_classes=10):
    super(Net, self).__init__()
    self.features = nn.Sequential(
      nn.Conv2d(1, 32, kernel_size=5, stride=1, padding=2),
      nn.MaxPool2d(kernel_size=2,stride=2),
      nn.Conv2d(32, 64, kernel_size=5, stride=1, padding=2),
      nn.MaxPool2d(kernel_size=2,stride=2),

    )
    self.classifier = nn.Sequential(
      nn.Linear(3136, 7*7*64),
      nn.Linear(3136, num_classes),

    )

  def forward(self,x):
    x = self.features(x)
    x = torch.flatten(x, 1)
    x = self.classifier(x)

    return x
net=Net()
net.cuda()#用GPU運行

#計算誤差，使用adam優(yōu)化器優(yōu)化誤差
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), 1e-2)

train_data = DataLoader(train_set, batch_size=128, shuffle=True)
test_data = DataLoader(test_set, batch_size=128, shuffle=False)


#訓(xùn)練過程
for epoch in range(1):
  net.train() ##在進(jìn)行訓(xùn)練時加上train()，測試時加上eval()
  batch = 0

  for batch_images, batch_labels in train_data:

    average_loss = 0
    train_acc = 0

    ##在pytorch0.4之后將Variable 與tensor進(jìn)行合并，所以這里不需要進(jìn)行Variable封裝
    if torch.cuda.is_available():
      batch_images, batch_labels = batch_images.cuda(),batch_labels.cuda()

    #前向傳播
    out = net(batch_images)
    loss = criterion(out,batch_labels)


    average_loss = loss
    prediction = torch.max(out,1)[1]
    # print(prediction)

    train_correct = (prediction == batch_labels).sum()
    ##這里得到的train_correct是一個longtensor型，需要轉(zhuǎn)換為float

    train_acc = (train_correct.float()) / 128

    optimizer.zero_grad() #清空梯度信息，否則在每次進(jìn)行反向傳播時都會累加
    loss.backward() #loss反向傳播
    optimizer.step() ##梯度更新

    batch+=1
    print("Epoch: %d/%d || batch:%d/%d average_loss: %.3f || train_acc: %.2f"
       %(epoch, 20, batch, float(int(50000/128)), average_loss, train_acc))

# 在測試集上檢驗效果
net.eval() # 將模型改為預(yù)測模式
for idx,(im1, label1) in enumerate(test_data):
  if torch.cuda.is_available():
    im, label = im1.cuda(),label1.cuda()
  out = net(im)
  loss = criterion(out, label)

  eval_loss = loss

  pred = torch.max(out,1)[1]
  num_correct = (pred == label).sum()
  acc = (num_correct.float())/ 128
  eval_acc = acc

  print('EVA_Batch:{}, Eval Loss: {:.6f}, Eval Acc: {:.6f}'
   .format(idx,eval_loss , eval_acc))