import numpy as np
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import time
np.random.seed(1234)

#這里用的是torchvision已經(jīng)封裝好的MINST數(shù)據(jù)集
trainset=torchvision.datasets.MNIST(
    root='MNIST',  #root是下載MNIST數(shù)據(jù)集保存的路徑，可以自行修改
    train=True,
    transform=torchvision.transforms.ToTensor(),
    download=True
)
 
testset=torchvision.datasets.MNIST(
    root='MNIST',
    train=False,
    transform=torchvision.transforms.ToTensor(),
    download=True
)
 
trainloader = DataLoader(dataset=trainset, batch_size=100, shuffle=True)   #DataLoader是一個很好地能夠幫助整理數(shù)據(jù)集的類，可以用來分批次，打亂以及多線程等操作
testloader = DataLoader(dataset=testset, batch_size=100, shuffle=True)

#可視化某一批數(shù)據(jù)
train_img,train_label=next(iter(trainloader))   #iter迭代器，可以用來便利trainloader里面每一個數(shù)據(jù)，這里只迭代一次來進(jìn)行可視化
fig, axes = plt.subplots(10, 10, figsize=(10, 10))
axes_list = []
#輸入到網(wǎng)絡(luò)的圖像
for i in range(axes.shape[0]):
    for j in range(axes.shape[1]):
        axes[i, j].imshow(train_img[i*10+j,0,:,:],cmap="gray")    #這里畫出來的就是我們想輸入到網(wǎng)絡(luò)里訓(xùn)練的圖像，與之對應(yīng)的標(biāo)簽用來進(jìn)行最后分類結(jié)果損失函數(shù)的計(jì)算
        axes[i, j].axis("off")
#對應(yīng)的標(biāo)簽
print(train_label)

#卷積模塊，由卷積核和激活函數(shù)組成
class conv_block(nn.Module):
    def __init__(self,ks,ch_in,ch_out):
        super(conv_block,self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(ch_in, ch_out, kernel_size=ks,stride=1,padding=1,bias=True),  #二維卷積核，用于提取局部的圖像信息
            nn.ReLU(inplace=True), #這里用ReLU作為激活函數(shù)
            nn.Conv2d(ch_out, ch_out, kernel_size=ks,stride=1,padding=1,bias=True),
            nn.ReLU(inplace=True),
        )
    def forward(self,x):
        return self.conv(x)

#常規(guī)CNN模塊（由幾個卷積模塊堆疊而成）
class CNN(nn.Module):
    def __init__(self,kernel_size,in_ch,out_ch):
        super(CNN, self).__init__()
        feature_list = [16,32,64,128,256]   #代表每一層網(wǎng)絡(luò)的特征數(shù)，擴(kuò)大特征空間有助于挖掘更多的局部信息
        self.conv1 = conv_block(kernel_size,in_ch,feature_list[0])
        self.conv2 = conv_block(kernel_size,feature_list[0],feature_list[1])
        self.conv3 = conv_block(kernel_size,feature_list[1],feature_list[2])
        self.conv4 = conv_block(kernel_size,feature_list[2],feature_list[3])
        self.conv5 = conv_block(kernel_size,feature_list[3],feature_list[4])
        self.fc =  nn.Sequential(           #全連接層主要用來進(jìn)行分類，整合采集的局部信息以及全局信息
            nn.Linear(feature_list[4] * 28 * 28, 1024),  #此處28為MINST一張圖片的維度
            nn.ReLU(),
            nn.Linear(1024, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )
 
    def forward(self,x):
        device = x.device
        x1 = self.conv1(x )
        x2 = self.conv2(x1)
        x3 = self.conv3(x2)
        x4 = self.conv4(x3)
        x5 = self.conv5(x4)
        x5 = x5.view(x5.size()[0], -1)  #全連接層相當(dāng)于做了矩陣乘法，所以這里需要將維度降維來實(shí)現(xiàn)矩陣的運(yùn)算
        out = self.fc(x5)
        return out

#網(wǎng)絡(luò)參數(shù)定義
device = torch.device("cuda:4")  #此處根據(jù)電腦配置進(jìn)行選擇，如果沒有cuda就用cpu
#device = torch.device("cpu")
net = CNN(3,1,1).to(device = device,dtype = torch.float32)
epochs = 50  #訓(xùn)練輪次
optimizer = torch.optim.Adam(net.parameters(), lr=1e-4, weight_decay=1e-8)  #使用Adam優(yōu)化器
criterion = nn.CrossEntropyLoss()  #分類任務(wù)常用的交叉熵?fù)p失函數(shù)
train_loss = []

# Begin training
MinTrainLoss = 999
for epoch in range(1,epochs+1):
    total_train_loss = []      
    net.train()
    start = time.time()
    for input_img,label in trainloader:
        input_img = input_img.to(device = device,dtype=torch.float32)  #我們同樣地，需要將我們?nèi)〕鰜淼挠?xùn)練集數(shù)據(jù)進(jìn)行torch能夠運(yùn)算的格式轉(zhuǎn)換
        label = label.to(device = device,dtype=torch.float32)          #輸入和輸出的格式都保持一致才能進(jìn)行運(yùn)算
        optimizer.zero_grad()  #每一次算loss前需要將之前的梯度清零，這樣才不會影響后面的更新
        pred_img = net(input_img) 
        loss = criterion(pred_img,label.long())
        loss.backward()
        optimizer.step()
        total_train_loss.append(loss.item())
    train_loss.append(np.mean(total_train_loss))    #將一個minibatch里面的損失取平均作為這一輪的loss
    end = time.time()
    #打印當(dāng)前的loss
    print("epochs[%3d/%3d] current loss: %.5f, time: %.3f"%(epoch,epochs,train_loss[-1],(end-start)))   #打印每一輪訓(xùn)練的結(jié)果
    
    if train_loss[-1]<MinTrainLoss:
        torch.save(net.state_dict(), "./model_min_train.pth")  #保存loss最小的模型
        MinTrainLoss = train_loss[-1]

#測試機(jī)某一批數(shù)據(jù)
test_img,test_label=next(iter(testloader))
fig, axes = plt.subplots(10, 10, figsize=(10, 10))
axes_list = []
#輸入到網(wǎng)絡(luò)的圖像
for i in range(axes.shape[0]):
    for j in range(axes.shape[1]):
        axes[i, j].imshow(test_img[i*10+j,0,:,:],cmap="gray")
        axes[i, j].axis("off")

#預(yù)測我拿出來的那一批數(shù)據(jù)進(jìn)行展示
cnn = CNN(3,1,1).to(device = device,dtype = torch.float32)
cnn.load_state_dict(torch.load("./model_min_train.pth", map_location=device)) #導(dǎo)入我們之前已經(jīng)訓(xùn)練好的模型
cnn.eval()   #評估模式
 
test_img = test_img.to(device = device,dtype = torch.float32)
test_label = test_label.to(device = device,dtype = torch.float32)
 
pred_test = cnn(test_img)  #記住，輸出的結(jié)果是一個長度為10的tensor
test_pred = np.argmax(pred_test.cpu().data.numpy(), axis=1)  #所以我們需要對其進(jìn)行最大值對應(yīng)索引的處理，從而得到我們想要的預(yù)測結(jié)果
 
#預(yù)測結(jié)果以及標(biāo)簽
print("預(yù)測結(jié)果")
print(test_pred)
print("標(biāo)簽")
print(test_label.cpu().data.numpy())