# 下載實(shí)驗(yàn)所需數(shù)據(jù)并解壓
!wget http://labfile.oss.aliyuncs.com/courses/1073/transfer-data.zip
!unzip transfer-data.zip

#引入實(shí)驗(yàn)所需要的包
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
import torch.nn.functional as F
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import copy
import os

# 數(shù)據(jù)存儲(chǔ)總路徑
data_dir = 'transfer-data'
# 圖像的大小為224*224
image_size = 224
# 從data_dir/train加載文件
# 加載的過程將會(huì)對(duì)圖像自動(dòng)作如下的圖像增強(qiáng)操作：
# 1. 隨機(jī)從原始圖像中切下來一塊224*224大小的區(qū)域
# 2. 隨機(jī)水平翻轉(zhuǎn)圖像
# 3. 將圖像的色彩數(shù)值標(biāo)準(zhǔn)化
train_dataset = datasets.ImageFolder(os.path.join(data_dir, 'train'),
                                    transforms.Compose([
                                        transforms.RandomResizedCrop(image_size),
                                        transforms.RandomHorizontalFlip(),
                                        transforms.ToTensor(),
                                        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                                    ])
                                    )

# 加載校驗(yàn)數(shù)據(jù)集，對(duì)每個(gè)加載的數(shù)據(jù)進(jìn)行如下處理：
# 1. 放大到256*256像素
# 2. 從中心區(qū)域切割下224*224大小的圖像區(qū)域
# 3. 將圖像的色彩數(shù)值標(biāo)準(zhǔn)化
val_dataset = datasets.ImageFolder(os.path.join(data_dir, 'val'),
                                    transforms.Compose([
                                        transforms.Resize(256),
                                        transforms.CenterCrop(image_size),
                                        transforms.ToTensor(),
                                        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                                    ])
                                    )

# 創(chuàng)建相應(yīng)的數(shù)據(jù)加載器
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size = 4, shuffle = True, num_workers=4)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size = 4, shuffle = True, num_workers=4)

# 讀取得出數(shù)據(jù)中的分類類別數(shù)
# 如果只有蜜蜂和螞蟻，那么是2
num_classes = len(train_dataset.classes)
num_classes

# 檢測(cè)本機(jī)器是否安裝GPU，將檢測(cè)結(jié)果記錄在布爾變量use_cuda中
use_cuda = torch.cuda.is_available()

# 當(dāng)可用GPU的時(shí)候，將新建立的張量自動(dòng)加載到GPU中
dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor
itype = torch.cuda.LongTensor if use_cuda else torch.LongTensor

def imshow(inp, title=None):
    # 將一張圖打印顯示出來，inp為一個(gè)張量，title為顯示在圖像上的文字

    # 一般的張量格式為：channels * image_width * image_height
    # 而一般的圖像為 image_width * image_height * channels 
    # 所以，需要將張量中的 channels 轉(zhuǎn)換到最后一個(gè)維度
    inp = inp.numpy().transpose((1, 2, 0)) 

    #由于在讀入圖像的時(shí)候所有圖像的色彩都標(biāo)準(zhǔn)化了，因此我們需要先調(diào)回去
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    inp = std * inp + mean
    inp = np.clip(inp, 0, 1) 

    #將圖像繪制出來
    plt.imshow(inp)
    if title is not None:
        plt.title(title)
    plt.pause(0.001)  # 暫停一會(huì)是為了能夠?qū)D像顯示出來。

#獲取第一個(gè)圖像batch和標(biāo)簽
images, labels = next(iter(train_loader))

# 將這個(gè)batch中的圖像制成表格繪制出來
out = torchvision.utils.make_grid(images)

imshow(out, title=[train_dataset.classes[x] for x in labels])

torch.utils.model_zoo.load_url('http://labfile.oss.aliyuncs.com/courses/1073/resnet18-5c106cde.pth')
# 加載模型庫中的residual network，并設(shè)置pretrained為true，這樣便可加載相應(yīng)的權(quán)重
net = models.resnet18(pretrained=True)
#如果存在GPU，就將網(wǎng)絡(luò)加載到GPU上
net = net.cuda() if use_cuda else net
# 將網(wǎng)絡(luò)的架構(gòu)打印出來
net

# 讀取最后線性層的輸入單元數(shù)，這是前面各層卷積提取到的特征數(shù)量
num_ftrs = net.fc.in_features

# 重新定義一個(gè)全新的線性層，它的輸出為2，原本是1000
net.fc = nn.Linear(num_ftrs, 2)

#如果存在GPU則將網(wǎng)絡(luò)加載到GPU中
net.fc = net.fc.cuda() if use_cuda else net.fc

criterion = nn.CrossEntropyLoss() #Loss函數(shù)的定義
# 將網(wǎng)絡(luò)的所有參數(shù)放入優(yōu)化器中
optimizer = optim.SGD(net.parameters(), lr = 0.0001, momentum=0.9)

record = [] #記錄準(zhǔn)確率等數(shù)值的容器

#開始訓(xùn)練循環(huán)
num_epochs = 20
net.train(True) # 給網(wǎng)絡(luò)模型做標(biāo)記，標(biāo)志說模型在訓(xùn)練集上訓(xùn)練
best_model = net
best_r = 0.0
for epoch in range(num_epochs):
    #optimizer = exp_lr_scheduler(optimizer, epoch)
    train_rights = [] #記錄訓(xùn)練數(shù)據(jù)集準(zhǔn)確率的容器
    train_losses = []
    for batch_idx, (data, target) in enumerate(train_loader):  #針對(duì)容器中的每一個(gè)批進(jìn)行循環(huán)
        data, target = Variable(data), Variable(target) #將Tensor轉(zhuǎn)化為Variable，data為圖像，target為標(biāo)簽
        #如果存在GPU則將變量加載到GPU中
        if use_cuda:
            data, target = data.cuda(), target.cuda()
        output = net(data) #完成一次預(yù)測(cè)
        loss = criterion(output, target) #計(jì)算誤差
        optimizer.zero_grad() #清空梯度
        loss.backward() #反向傳播
        optimizer.step() #一步隨機(jī)梯度下降
        right = rightness(output, target) #計(jì)算準(zhǔn)確率所需數(shù)值，返回正確的數(shù)值為（正確樣例數(shù)，總樣本數(shù)）
        train_rights.append(right) #將計(jì)算結(jié)果裝到列表容器中
        loss = loss.cpu() if use_cuda else loss
        train_losses.append(loss.data.numpy())


        #if batch_idx % 20 == 0: #每間隔100個(gè)batch執(zhí)行一次
     #train_r為一個(gè)二元組，分別記錄訓(xùn)練集中分類正確的數(shù)量和該集合中總的樣本數(shù)
    train_r = (sum([tup[0] for tup in train_rights]), sum([tup[1] for tup in train_rights]))

    #在測(cè)試集上分批運(yùn)行，并計(jì)算總的正確率
    net.eval() #標(biāo)志模型當(dāng)前為運(yùn)行階段
    test_loss = 0
    correct = 0
    vals = []
    #對(duì)測(cè)試數(shù)據(jù)集進(jìn)行循環(huán)
    for data, target in val_loader:
        #如果存在GPU則將變量加載到GPU中
        if use_cuda:
            data, target = data.cuda(), target.cuda()
        data, target = Variable(data, requires_grad=True), Variable(target)
        output = net(data) #將特征數(shù)據(jù)喂入網(wǎng)絡(luò)，得到分類的輸出
        val = rightness(output, target) #獲得正確樣本數(shù)以及總樣本數(shù)
        vals.append(val) #記錄結(jié)果

    #計(jì)算準(zhǔn)確率
    val_r = (sum([tup[0] for tup in vals]), sum([tup[1] for tup in vals]))
    val_ratio = 1.0*val_r[0].numpy()/val_r[1]

    if val_ratio > best_r:
        best_r = val_ratio
        best_model = copy.deepcopy(net)
    #打印準(zhǔn)確率等數(shù)值，其中正確率為本訓(xùn)練周期Epoch開始后到目前撮的正確率的平均值
    print('訓(xùn)練周期: {} \tLoss: {:.6f}\t訓(xùn)練正確率: {:.2f}%, 校驗(yàn)正確率: {:.2f}%'.format(
        epoch, np.mean(train_losses), 100. * train_r[0].numpy() / train_r[1], 100. * val_r[0].numpy()/val_r[1]))       
    record.append([np.mean(train_losses), 1. * train_r[0].data.numpy() / train_r[1], 1. * val_r[0].data.numpy() / val_r[1]])

#繪制訓(xùn)練誤差曲線
x = [x[0] for x in record]
y = [1 - x[1] for x in record]
z = [1 - x[2] for x in record]
#plt.plot(x)
plt.figure(figsize = (10, 7))
plt.plot(y)
plt.plot(z)
plt.xlabel('Epoch')
plt.ylabel('Error Rate')

def visualize_model(model, num_images=6):
    images_so_far = 0
    fig = plt.figure(figsize=(15,10))

    for i, data in enumerate(val_loader):
        inputs, labels = data
        inputs, labels = Variable(inputs), Variable(labels)
        if use_cuda:
            inputs, labels = inputs.cuda(), labels.cuda()
        outputs = model(inputs)
        _, preds = torch.max(outputs.data, 1)
        preds = preds.cpu().numpy() if use_cuda else preds.numpy()
        for j in range(inputs.size()[0]):
            images_so_far += 1
            ax = plt.subplot( 2,num_images//2, images_so_far)
            ax.axis('off')

            ax.set_title('predicted: {}'.format(val_dataset.classes[preds[j]]))
            imshow(data[0][j])

            if images_so_far == num_images:
                return
visualize_model(net)

plt.ioff()
plt.show()

# 加載residual網(wǎng)絡(luò)模型
net = torchvision.models.resnet18(pretrained=True)
# 將模型放入GPU中
net = net.cuda() if use_cuda else net

# 循環(huán)網(wǎng)絡(luò)，將所有參數(shù)設(shè)為不更新梯度信息
for param in net.parameters():
    param.requires_grad = False

# 將網(wǎng)絡(luò)最后一層線性層換掉
num_ftrs = net.fc.in_features
net.fc = nn.Linear(num_ftrs, 2)
net.fc = net.fc.cuda() if use_cuda else net.fc

criterion = nn.CrossEntropyLoss() #Loss函數(shù)的定義
# 僅將線性層的參數(shù)放入優(yōu)化器中
optimizer = optim.SGD(net.fc.parameters(), lr = 0.001, momentum=0.9)


#訓(xùn)練模型
record = [] #記錄準(zhǔn)確率等數(shù)值的容器

#開始訓(xùn)練循環(huán)
num_epochs = 4
net.train(True) # 給網(wǎng)絡(luò)模型做標(biāo)記，標(biāo)志說模型在訓(xùn)練集上訓(xùn)練
best_model = net
best_r = 0.0
for epoch in range(num_epochs):
    #optimizer = exp_lr_scheduler(optimizer, epoch)
    train_rights = [] #記錄訓(xùn)練數(shù)據(jù)集準(zhǔn)確率的容器
    train_losses = []
    for batch_idx, (data, target) in enumerate(train_loader):  #針對(duì)容器中的每一個(gè)批進(jìn)行循環(huán)
        data, target = Variable(data), Variable(target) #將Tensor轉(zhuǎn)化為Variable，data為圖像，target為標(biāo)簽
        if use_cuda:
            data, target = data.cuda(), target.cuda()
        output = net(data) #完成一次預(yù)測(cè)
        loss = criterion(output, target) #計(jì)算誤差
        optimizer.zero_grad() #清空梯度
        loss.backward() #反向傳播
        optimizer.step() #一步隨機(jī)梯度下降
        right = rightness(output, target) #計(jì)算準(zhǔn)確率所需數(shù)值，返回正確的數(shù)值為（正確樣例數(shù)，總樣本數(shù)）
        train_rights.append(right) #將計(jì)算結(jié)果裝到列表容器中
        loss = loss.cpu() if use_cuda else loss
        train_losses.append(loss.data.numpy())


     #train_r為一個(gè)二元組，分別記錄訓(xùn)練集中分類正確的數(shù)量和該集合中總的樣本數(shù)
    train_r = (sum([tup[0] for tup in train_rights]), sum([tup[1] for tup in train_rights]))

    #在測(cè)試集上分批運(yùn)行，并計(jì)算總的正確率
    net.eval() #標(biāo)志模型當(dāng)前為運(yùn)行階段
    test_loss = 0
    correct = 0
    vals = []
    #對(duì)測(cè)試數(shù)據(jù)集進(jìn)行循環(huán)
    for data, target in val_loader:
        data, target = Variable(data, requires_grad=True), Variable(target)
        if use_cuda:
            data, target = data.cuda(), target.cuda()
        output = net(data) #將特征數(shù)據(jù)喂入網(wǎng)絡(luò)，得到分類的輸出
        val = rightness(output, target) #獲得正確樣本數(shù)以及總樣本數(shù)
        vals.append(val) #記錄結(jié)果

    #計(jì)算準(zhǔn)確率
    val_r = (sum([tup[0] for tup in vals]), sum([tup[1] for tup in vals]))
    val_ratio = 1.0*val_r[0].numpy()/val_r[1]

    if val_ratio > best_r:
        best_r = val_ratio
        best_model = copy.deepcopy(net)
    #打印準(zhǔn)確率等數(shù)值，其中正確率為本訓(xùn)練周期Epoch開始后到目前撮的正確率的平均值
    print('訓(xùn)練周期: {} \tLoss: {:.6f}\t訓(xùn)練正確率: {:.2f}%, 校驗(yàn)正確率: {:.2f}%'.format(
        epoch, np.mean(train_losses), 100. * train_r[0].numpy() / train_r[1], 100. * val_r[0].numpy()/val_r[1]))       
    record.append([np.mean(train_losses), 1. * train_r[0].data.numpy() / train_r[1], 1. * val_r[0].data.numpy() / val_r[1]])


# 繪制誤差曲線
x = [x[0] for x in record]
y = [1 - x[1] for x in record]
z = [1 - x[2] for x in record]
#plt.plot(x)
plt.figure(figsize = (10, 7))
plt.plot(y)
plt.plot(z)
plt.xlabel('Epoch')
plt.ylabel('Error Rate')


#展示分類結(jié)果
visualize_model(best_model)

plt.ioff()
plt.show()