【PyTorch】圖像分類"貓狗大戰"實戰

項目總體目錄

├── pytorch_dogsVScats
│   ├── datas
│   │   ├── Dataset
│   │   ├── train
│   │   │   ├── dog
│   │   │   └── cat
│   │   └── valid
│   │       ├── dog
│   │       └── cat
│   ├── data
│   │   └── dataset.py
│   │
│   ├── models 
│   │   ├── VGG.py
│   │   └── AlexNet.py
│   │
│   ├── utils 
│   │   └── get_mean_std.py
│   │
│   ├── config.py 
│   ├── redistributionImgs.py
│   └── main.py

1. 數據集下載和預覽

1.1 數據集下載

到Kaggle官網下載貓狗分類的數據集。該數據集的訓練集共有25000張貓和狗的圖片，測試集共有12500張貓和狗的圖片。

在實驗中，我們不會直接使用測試數據集對搭建的模型進行訓練和優化，而是在訓練數據集中劃出 一部分作爲驗證集，來評估在每個批次的訓練後模型的泛化能力。這樣做的 原因是如果我們使用測試數據集進行模型訓練和優化，那麼模型最終會對測試數據集產生擬合傾向，換而言之，我們的模型只有在對測試數據集中圖片的類別進行預測時纔有極強 的準確率，而在對測試數據集以外的圖片類別進行預測時會出現非常多的錯誤，這樣的模型缺少泛化能力。所以，爲了防止這種情況的出現，我們會把測試數據集從模型的訓練和優化過程中隔離出來，只在每輪訓練結束後使用。如果模型對驗證數據集和測試數據集的 預測同時具備高準確率和l低損失值，就基本說明模型的參數優化是成功的 ， 模型將具備極強的泛化能力。

新建項目名爲pytorch_dogsVScats，並在項目根目錄下創建如下所示的目錄結構：

├── datas
│   ├── Dataset
│   ├── train
│   │   ├── dog
│   │   └── cat
│   └── valid
│   │   ├── dog
│   │   └── cat

datas爲項目存放數據的根目錄，將下載的所有圖片共都放置到Dataset目錄下，train和valid目錄分別用於存儲訓練圖片和驗證圖片，本次實驗將訓練集中20%的圖片（貓狗分別2500張，共5000張）用來組成驗證數據集。

新建redistributionImgs.py文件，編寫redistributionImgs函數來將圖片分別移動到相應的目錄下。具體實現如下：

import os
import shutil
from config import Config

def redistributionImgs():
    '''
    redistributing all pictures(in './data/Dataset') to the corresponding directory
    '''

    print('Redistribution start...')
    # 調用配置文件
    conf = Config()
    # 如果圖片還未分配則進行分配
    if not os.listdir(os.path.join(conf.data_train_root, 'cat/')):
        # 返回所有圖片的名字列表
        file_datas = os.listdir(os.path.join(conf.data_root, 'Dataset/'))
        # 使用filter()從有標籤的25000張圖片中過濾圖片
        # 使用匿名函數lambda處理篩選條件
        # 返回標籤爲狗的圖片
        file_dogs = list(filter(lambda x: x[:3] == 'dog', file_datas))
        # 返回標籤爲貓的圖片
        file_cats = list(filter(lambda x: x[:3] == 'cat', file_datas))
        # 有標籤的狗圖和貓圖的個數
        d_len, c_len = len(file_dogs), len(file_cats)
        # 80%的圖片用於訓練，20%的圖片用於驗證
        val_d_len, val_c_len = d_len * 0.8, c_len * 0.8

        # 分配狗圖片
        for i in range(d_len):
            pre_path = os.path.join(conf.data_root, 'Dataset', file_dogs[i])
            # 80%的數據作爲訓練數據集
            if i < val_d_len:
                new_path = os.path.join(conf.data_train_root, 'dog/')
            # 20%的數據作爲驗證數據集
            else:
                new_path = os.path.join(conf.data_valid_root, 'dog/')
            # 調用shutil.move()移動文件
            shutil.move(pre_path, new_path)

        # 分配貓圖片
        for i in range(c_len):
            pre_path = os.path.join(conf.data_root, 'Dataset', file_cats[i])
            # 80%的數據作爲訓練數據集
            if i < val_c_len:
                new_path = os.path.join(conf.data_train_root, 'cat/')
            # 20%的數據作爲驗證數據集
            else:
                new_path = os.path.join(conf.data_valid_root, 'cat/')
            # 調用shutil.move()移動文件
            shutil.move(pre_path, new_path)
        print('Redistribution completed!')

新建config.py文件用於配置實驗中的經常使用變量。

class Config:
    # 文件路徑
    # 數據集根目錄
    data_root = './datas/'
    # 訓練集存放路徑
    data_train_root = './datas/train/'
    # 驗證集存放路徑
    data_valid_root = './datas/valid/'
    # 測試集存放路徑
    data_test_root = './datas/test/'

    # 測試結果保存位置
    result_file = './result.csv'

    # 常用參數
    # batch size
    batch_size = 32
    # mean and std
    # 通過抽樣計算得到圖片的均值mean和標準差std
    mean = [0.470, 0.431, 0.393]
    std = [0.274, 0.263, 0.260]

    # 預訓練模型路徑
    path_vgg16 = './models/vgg16-397923af.pth'

get_mean_std用於計算原數據的均值和方差。

def get_mean_std(data_images):
    '''
    :param data_images: 加載好的數據集
    :return: mean,std
    '''
    times, mean, std = 0, 0, 0
    data_loader = {x: torch.utils.data.DataLoader(dataset=data_images[x],
                                                  batch_size=1000,
                                                  shuffle=True)
                          for x in ['train', 'valid']}
    for imgs, labels in data_loader['train']:
        # imgs.shape = torch.Size([32, 3, 64, 64])
        times += 1
        mean += np.mean(imgs.numpy(), axis=(0, 2, 3))
        std += np.std(imgs.numpy(), axis=(0, 2, 3))
        print('times:', times)

    mean /= times
    std /= times
    return mean, std

1.2 數據加載和預覽

在項目根目錄下新建data目錄，並在data目錄下新建dataset.py文件，用於實現數據的加載。

在代碼中對數據的變換和導入都使用了字典的形式，因爲我們需要分別對訓練數據集和驗證數據集的數據載入方法進行簡單定義，所以使用字典可以簡化代碼，也方便之後進行相應的調用和操作。

在加載完數據後，並對一個批次的數據進行預覽。
dataset.py中的代碼如下：

# 調用配置文件
conf = Config()

class Dataset:
    def __init__(self, train=True):
        # 圖片預處理
        # Compose用於將多個transfrom組合起來
        # ToTensor()將像素從[0, 255]轉換爲[0, 1.0]
        # Normalize()用均值和標準差對圖像標準化處理 x'=(x-mean)/std，加速收斂
        self.transform = transforms.Compose([transforms.Resize((64, 64)),
                                             transforms.ToTensor(),
                                             transforms.Normalize(conf.mean, conf.std)])
        self.train = train
        # 加載訓練數據集和驗證數據集
        if train:

            # 數據加載
            # 這裏使用通用的ImageFolder和DataLoader數據加載器
            # 數據類型 data_images = {'train': xxx, 'valid': xxx}
            self.data_images = {x: datasets.ImageFolder(root=os.path.join(conf.data_root, x),
                                                        transform=self.transform)
                                for x in ['train', 'valid']}
            self.data_images_loader = {x: torch.utils.data.DataLoader(dataset=self.data_images[x],
                                                                      batch_size=conf.batch_size,
                                                                      shuffle=True)
                                       for x in ['train', 'valid']}
            # 圖片分類 ['cat', 'dog']
            self.classes = self.data_images['train'].classes
            # 圖片分類鍵值對 {'cat': 0, 'dog': 1}
            self.classes_index = self.data_images['train'].class_to_idx

        # 加載測試數據集
        else:
            images = [os.path.join(conf.data_test_root, img) for img in os.listdir(conf.data_test_root)]
            self.images = sorted(images, key=lambda x: int(x.split('.')[-2].split('/')[-1]))

    # 重載專有方法__getitem__
    def __getitem__(self, index):
        img_path = self.images[index]
        label = int(self.images[index].split('.')[-2].split('/')[-1])
        data_images_test = Image.open(img_path)
        data_images_test = self.transform(data_images_test)
        return data_images_test, label

    # 重載專有方法__len__
    def __len__(self):
        return len(self.images)

main.py中的代碼：

from data.dataset import Dataset

dst = Dataset()
# 預覽其中一個批次的數據
imgs, labels = iter(dst.data_images_loader['train']).next()
print(imgs.shape)
print(labels)
print(labels.shape)

打印結果:

torch.Size([32, 3, 224, 224])
tensor([0, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 1, 1, 0, 0,
        1, 0, 0, 1, 1, 0, 1, 0])
torch.Size([32])

imgs和labels都是Tensor數據類型的變量。因爲我們對圖片的大小進行了縮放變換，所以圖片現在的大小都是64 × 64，因此imgs的維度就是(32, 3, 64, 64)，32表示這個批次中有32張圖片；3表示色彩通道數，因爲原始圖像是彩色的，因此使用了R、G、B三個通道；64表示圖片的寬和高。

labels中的元素全部是0和1，因爲在進行數據裝載時已經對dog文件夾和cat文件夾下的內容進行了獨熱編碼（One-Hot-Encoding），因此這時的0和1不僅是每張圖片的標籤，還分別對應貓的圖片和狗的圖片。這裏做一個簡單的打印輸出，來驗證這個獨熱編碼的對應關係，代碼如下：

classes_index = data_images['train'].class_to_idx
print(classes_index)

打印結果：

{'cat': 0, 'dog': 1}

相較於使用文字作爲圖片的標籤而言，使用 0和l也可以讓之後的計算方便很多。

爲了增加之後繪製的圖像標籤的可識別性，將原始標籤的結果存儲在名爲classes的變量中。

classes = data_images['train'].classes
print(classes)

輸出結果：

['cat', 'dog']

使用opencv的cv2庫來預覽圖片，代碼如下：

# 圖片預覽
    imgs, labels = iter(dst.data_images_loader['train']).next()
    # 製作雪碧圖
    # 類型爲tensor，維度爲[channel, height, width]
    img = torchvision.utils.make_grid(imgs)
    # 轉換爲數組並調整維度爲[height, width, channel]
    img = img.numpy().transpose([1, 2, 0])
    # 通過反向推導標準差交換法計算圖片原來的像素值
    mean, std = conf.mean, conf.std
    img = img * std + mean
    # 打印圖片標籤
    print([dst.classes[i] for i in labels])
    # 顯示圖片
    cv2.imshow('img', img)
    # 等待圖片關閉
    cv2.waitKey(0)

輸出結果：

['cat', 'dog', 'cat', 'cat', 'cat', 'cat', 'dog', 'cat', 'dog', 'dog', 'dog', 'dog', 'dog', 'cat', 'cat', 'dog', 'dog', 'dog', 'cat', 'cat', 'cat', 'cat', 'dog', 'cat', 'dog', 'dog', 'cat', 'dog', 'dog', 'cat', 'dog', 'dog']

2. 搭建網絡模型

我們開始搭建卷積神經網絡模型，考慮到硬件環境和訓練花費的時間成本，本次實驗我們採用簡化版的VGG-16網絡模型。標準的VGG-16模型中要求的輸入圖片大小爲224×224，模型中共有有13層卷積層和5層池化層。我們的模型將輸入的圖片大小全部縮放爲64×64，同時刪除了原模型最後的3個卷積層和池化層，並且改變了全連接層中的連接參數，原模型中全連接層的輸入維度爲7×7×512，因爲我們的輸入圖片大小爲64×64，因此通過計算得出的全連接層的輸入維度爲4×4×512。通過調整模型的大小和參數，大幅度減少了整個模型參與訓練的參數數量。

簡化後的VGG模型的代碼如下：

import torch.nn as nn

# 需繼承torch.nn.Module類
class VGG16(nn.Module):

    def __init__(self):
        super(VGG16, self).__init__()
        # 定義卷積層和池化層，共13層卷積，5層池化
        self.conv = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),

            # nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            # nn.ReLU(),
            # nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            # nn.ReLU(),
            # nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            # nn.ReLU(),
            # nn.MaxPool2d(kernel_size=2, stride=2),
        )

        # 簡化版全連接層
        self.classes = nn.Sequential(
            nn.Linear(4 * 4 * 512, 1024),
            nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(1024, 1024),
            nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(1024, 2)
        )

        # VGG-16的全連接層
        # self.classes = nn.Sequential(
        #     nn.Linear(7 *7 *512, 4096),
        #     nn.ReLU(),
        #     nn.Dropout(p=0.5),
        #     nn.Linear(4096, 4096),
        #     nn.ReLU(),
        #     nn.Dropout(p=0.5),
        #     nn.Linear(4096, 2)
        # )

    # 定義每次執行的計算步驟
    def forward(self, x):
        x = self.conv(x)
        x = x.view(-1, 4 * 4* 512)
        x = self.classes(x)
        return x

3. 模型訓練

定義模型的損失函數和對參數進行優化的優化函數，這裏損失函數採用交叉熵torch.nn.CrossEntropyLoss，優化函數採用torch.optim.Adam。代碼如下：

# 定義損失函數
loss_fn = nn.CrossEntropyLoss()
# 定義優化器，優化模型上的所有參數和學習率，默認lr=1e-3
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

將在模型訓練的過程中需要計算的參數全部遷移至GPUs上，這個過程非常簡單和方便，只需重新對這部分參數進行類型轉換就可以了。但在啓用GUPs前需要檢查GPUs硬件是否可用，代碼如下：

# 如果GPUs可用，則將模型上需要計算的所有參數複製到GPUs上
if torch.cuda.is_available():
    model = model.cuda()

將需要計算的參數遷移至GPUs上，代碼如下：

if torch.cuda.is_available():
    X, y = X.cuda(), y.cuda()

main.py中的完整代碼如下：

# 數據類實例
dst = Dataset()
# 模型類實例
model = VGG16()
# 配置類實例
conf = Config()

# 定義損失函數
loss_fn = nn.CrossEntropyLoss()
# 定義優化器，優化模型上的所有參數和學習率，默認lr=1e-3
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# 定義超級參數
epoch_n = 10

# 如果GPUs可用，則將模型上需要計算的所有參數複製到GPUs上
if torch.cuda.is_available():
    model = model.cuda()

# 代碼執行時間裝飾器
def timer(func):
    def wrapper(*args, **kw):
        begin = time.time()
        # 執行函數體
        func(*args, **kw)
        end = time.time()

        # 花費時間
        cost = end - begin
        print('本次一共花費時間：{:.2f}秒。'.format(cost))
    return wrapper


# 訓練
@timer
def train():
    for epoch in range(1, epoch_n + 1):
        print('Epoch {}/{}'.format(epoch, epoch_n))
        print('-'*20)

        for phase in ['train', 'valid']:
            if phase == 'train':
                print('Training...')
                # 打開訓練模式
                model.train(True)
            else:
                print('Validing...')
                # 關閉訓練模式
                model.train(False)

            # 損失值
            running_loss = 0.0
            # 預測的正確數
            running_correct = 0
            # 讓batch的值從1開始，便於後面計算
            for batch, data in enumerate(dst.data_images_loader[phase], 1):
                # 實際輸入值和輸出值
                X, y = data
                # 將參數複製到GPUs上進行運算
                if torch.cuda.is_available():
                    X, y = X.cuda(), y.cuda()
                # outputs.shape = [32,2] -> [1,2]
                outputs = model(X)
                # 從輸出結果中取出需要的預測值
                # axis = 1 表示去按行取最大值
                _, y_pred = torch.max(outputs.detach(), 1)
                # 將累積的梯度置零
                optimizer.zero_grad()
                # 計算損失值
                loss = loss_fn(outputs, y)
                if phase == 'train':
                    # 反向傳播求導
                    loss.backward()
                    # 更新所有參數
                    optimizer.step()

                running_loss += loss.detach().item()
                running_correct += torch.sum(y_pred == y)

                if batch % 500 == 0 and phase == 'train':
                    print('Batch {}/{},Train Loss:{:.2f},Train Acc:{:.2f}%'.format(
                        batch, len(dst.data_images[phase])/conf.batch_size, running_loss/batch, 100*running_correct.item()/(conf.batch_size*batch)
                    ))
            epoch_loss = running_loss*conf.batch_size/len(dst.data_images[phase])
            epoch_acc = 100*running_correct.item()/len(dst.data_images[phase])
            print('{} Loss:{:.2f} Acc:{:.2f}%'.format(phase, epoch_loss, epoch_acc))

通過10次訓練後，打印輸出如下：

Epoch 1/10
--------------------
Training...
Batch 500/625.0,Train Loss:0.69,Train Acc:50.52%
train Loss:0.69 Acc:51.07%
Validing...
valid Loss:0.70 Acc:50.12%
Epoch 2/10
--------------------
Training...
Batch 500/625.0,Train Loss:0.69,Train Acc:51.04%
train Loss:0.69 Acc:51.74%
Validing...
valid Loss:0.72 Acc:58.90%
Epoch 3/10
--------------------
Training...
Batch 500/625.0,Train Loss:0.63,Train Acc:63.59%
train Loss:0.62 Acc:64.86%
Validing...
valid Loss:0.56 Acc:72.38%
Epoch 4/10
--------------------
Training...
Batch 500/625.0,Train Loss:0.53,Train Acc:73.36%
train Loss:0.53 Acc:73.89%
Validing...
valid Loss:0.48 Acc:77.00%
Epoch 5/10
--------------------
Training...
Batch 500/625.0,Train Loss:0.46,Train Acc:79.17%
train Loss:0.45 Acc:79.68%
Validing...
valid Loss:0.40 Acc:81.36%
Epoch 6/10
--------------------
Training...
Batch 500/625.0,Train Loss:0.38,Train Acc:83.01%
train Loss:0.38 Acc:83.14%
Validing...
valid Loss:0.35 Acc:84.24%
Epoch 7/10
--------------------
Training...
Batch 500/625.0,Train Loss:0.33,Train Acc:85.78%
train Loss:0.32 Acc:86.11%
Validing...
valid Loss:0.38 Acc:83.56%
Epoch 8/10
--------------------
Training...
Batch 500/625.0,Train Loss:0.27,Train Acc:88.37%
train Loss:0.27 Acc:88.36%
Validing...
valid Loss:0.33 Acc:86.54%
Epoch 9/10
--------------------
Training...
Batch 500/625.0,Train Loss:0.23,Train Acc:90.61%
train Loss:0.23 Acc:90.61%
Validing...
valid Loss:0.29 Acc:87.92%
Epoch 10/10
--------------------
Training...
Batch 500/625.0,Train Loss:0.19,Train Acc:92.36%
train Loss:0.19 Acc:92.41%
Validing...
valid Loss:0.26 Acc:89.16%
本次一共花費時間：947.35秒。

4. 模型測試

# 測試
def test():
    dst = Dataset(train=False)
    data_loader_test = torch.utils.data.DataLoader(dst,
                                                   batch_size=conf.batch_size,
                                                   shuffle=False)
    # 保存測試結果
    results = []
    # tqdm模塊用於顯示進度條
    for imgs, path in tqdm(data_loader_test):
        if torch.cuda.is_available():
            X = imgs.cuda()
        outputs = model(X)
        # pred表示是哪個對象，0=cat，1=dog
        # probability表示是否個對象的概率
        probability, pred = torch.max(F.softmax(outputs, dim=1).detach(), 1)
        # 通過zip()打包爲元組的列表,如[(1001, 23%, 'cat')]
        batch_results = [(path_.item(), str(round(probability_.item()*100, 2))+'%', 'dog'if pred_.item() else 'cat')
                         for path_, probability_, pred_ in zip(path, probability, pred)]
        results += batch_results
    write_csv(results, conf.result_file)

    # 隨機預覽
    # 預覽圖片張數
    size = 20
    for i in range(size):
        index = math.floor(random.random() * len(results))
        # 圖片名稱
        file_name = str(results[index][0]) + '.jpg'
        # 圖片是否個對象的概率
        score = results[index][1]
        # 'cat' or 'dog'
        name = results[index][2]
        # 獲取圖片
        img = cv2.imread(os.path.join(conf.data_test_root, file_name))
        title = score + ' is ' + name

        # 設置圖片大小
        plt.rcParams['figure.figsize'] = (8.0, 8.0)
        # 默認間距
        plt.tight_layout()
        # 行，列，索引
        plt.subplot(4, 5, i + 1)
        plt.imshow(img)
        plt.title(title, fontsize=14, color='blue')
        plt.xticks([])
        plt.yticks([])
    plt.show()

# 保存測試結果
def write_csv(results, file_name):
    import csv
    with open(file_name, 'w') as f:
        writer = csv.writer(f)
        writer.writerow(['id', 'label', 'pred'])
        writer.writerows(results)

測試結果：

100%|██████████| 391/391 [00:45<00:00,  8.55it/s]

遷移學習-遷移VGG16

遷移學習其基本思路是凍結卷積神經網絡中全連接層之前的全部網絡層次，讓這些被凍結的網絡層次中的參數在模型的訓練過程中不進行梯度更新，能夠被優化的參數僅僅是沒有被凍結的全連接層的全部參數。

下面看看遷移學習的實施過程，首先需要下載己經具備最優參數的模型，這需要 對我們之前使用的 model = Models()代碼部分進行替換，因爲我們不需要再自己搭建和定義訓練的模型了，而是通過代碼自動下載模型並直接調用，具體代碼如下 ：

# pretrained=True表示下載訓練好的模型
# 由於下載速度太慢，這裏設置爲False
model = models.vgg16(pretrained=False)

# 加載預先下載好的預訓練參數到vgg16
model.load_state_dict(torch.load(conf.path_vgg16))

for param in model.parameters():
    param.requires_grad = False

model.classifier = nn.Sequential(
    nn.Linear(7*7*512, 4096),
    nn.ReLU(),
    nn.Dropout(p=0.5),
    nn.Linear(4096, 4096),
    nn.ReLU(),
    nn.Dropout(p=0.5),
    nn.Linear(4096, 2)
)

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.classifier.parameters(), lr=1e-5)

首先，遷移過來的VGG16架構模型在最後輸出的結果是1000個，在我們的問題中只需兩個輸出結果，所以全連接層必須進行調整。

其次，對原模型中的參數進行遍歷操作，將參數中的parma.requires_grad全部設置爲 False，這樣對應的參數將不計算梯度，當然也不會進行梯度更新了，這就是之前說到的凍結操作。然後，定義新的全連接層結構並重新賦值給model.classifier。在完成了新的全連接層定義後，全連接層中的parma.requires_grad參數會被默認重置爲True，所以不需要再次遍歷參數來進行解凍操作。損失函數的loss值依然使用交叉熵進行計算 ，但是在優化函數中負責優化的參數變成了全連接層中的所有參數， 即對 model.classifier.parameters 這部分參數進行優化。

問題及錯誤處理

1.1 CUDA out of memory

從問題中看出是因爲GPU的內存不足。通過查閱資料找到了一些解決方案：

1. 修改batch_size的大小
2. 簡化訓練模型
3. 釋放你不需要的張量和變量
4. 修改輸入圖片的大小

1.2 numpy中axis的取值問題

官方文檔中描述：

axis : None或int或tuple of ints，可選

沿一個或多個軸執行邏輯AND規約。 默認值（axis = None）是對輸入數組的所有維度執行邏輯AND。軸可以是負數，這種情況下，它從最後一個軸開始索引。

版本1.7.0中的新功能。

如果這是一個整數元組，則在多個軸上進行規約操作，而不是像以前那樣在單個軸或所有軸上進行行規約操作。

在本實驗中我們需要求出R、G、B三個通道上像素的均值和標準差。

因爲傳入的圖像維度爲[16, 3, 64, 64]，其中16表示一個批次的16張圖片，3表示通道數爲3，兩個64分別表示圖片的寬和高的大小。
因此在求三個通道的像素的均值和標準差時傳入的axis=(0, 2, 3)。代碼如下：

for imgs, labels in data_loader['train']:
    # imgs.shape = torch.Size([16, 3, 64, 64])
    mean += np.mean(imgs.numpy(), axis=(0, 2, 3))
    std += np.std(imgs.numpy(), axis=(0, 2, 3))

查看項目完整代碼請移駕至github。