目錄

1. PyTorch 預(yù)訓(xùn)練模型

2. 保存模型參數(shù)

3. 讀取模型參數(shù)

4. 凍結(jié)部分模型參數(shù)，進(jìn)行 fine-tuning

5. 模型訓(xùn)練與測試的設(shè)置

6. 利用 torch.nn.DataParallel 進(jìn)行多 GPU 訓(xùn)練

1. PyTorch 預(yù)訓(xùn)練模型

Pytorch 提供了許多 Pre-Trained Model on ImageNet，僅需調(diào)用 torchvision.models 即可，具體細(xì)節(jié)可查看官方文檔。

往往我們需要對 Pre-Trained Model 進(jìn)行相應(yīng)的修改，以適應(yīng)我們的任務(wù)。這種情況下，我們可以先輸出 Pre-Trained Model 的結(jié)構(gòu)，確定好對哪些層修改，或者添加哪些層，接著，再將其修改即可。

比如，我需要將 ResNet-50 的 Layer 3 后的所有層去掉，在分別連接十個(gè)分類器，分類器由 ResNet-50.layer4 和 AvgPool Layer 和 FC Layer 構(gòu)成。這里就需要用到 torch.nn.ModuleList 了，比如：：self.linears = nn.ModuleList([nn.Linear(10, 10) for i in range(10)])。

代碼中的 [nn.Linear(10, 10) for i in range(10)] 是一個(gè)python列表，必須要把它轉(zhuǎn)換成一個(gè)Module Llist列表才可以被 PyTorch 使用，否則在運(yùn)行的時(shí)候會(huì)報(bào)錯(cuò): RuntimeError: Input type (CUDAFloatTensor) and weight type (CPUFloatTensor) should be the same

2. 保存模型參數(shù)

PyTorch 中保存模型的方式有許多種：

3. 讀取模型參數(shù)

同樣的，PyTorch 中讀取模型參數(shù)的方式也有許多種：

# 讀取整個(gè)網(wǎng)絡(luò)model = torch.load(PATH)# 讀取網(wǎng)絡(luò)中的參數(shù)model.load_state_dict(torch.load(PATH))# 讀取網(wǎng)絡(luò)中的部分參數(shù)（本質(zhì)其實(shí)就是更新字典）pretrained_dict = torch.load(pretrained_model_weight)model_dict = model.state_dict()pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}model_dict.update(pretrained_dict)

4. 凍結(jié)部分模型參數(shù)，進(jìn)行 fine-tuning

加載完 Pre-Trained Model 后，我們需要對其進(jìn)行 Finetune。但是在此之前，我們往往需要凍結(jié)一部分的模型參數(shù)：

5. 模型訓(xùn)練與測試的設(shè)置

訓(xùn)練時(shí)，應(yīng)調(diào)用 model.train() ；測試時(shí)，應(yīng)調(diào)用 model.eval()，以及 with torch.no_grad():

model.train()：使 model 變成訓(xùn)練模式，此時(shí) dropout 和 batch normalization 的操作在訓(xùn)練起到防止網(wǎng)絡(luò)過擬合的問題。

model.eval()：PyTorch會(huì)自動(dòng)把 BN 和 DropOut 固定住，不會(huì)取平均，而是用訓(xùn)練好的值。不然的話，一旦測試集的 Batch Size 過小，很容易就會(huì)被 BN 層導(dǎo)致生成圖片顏色失真極大。

with torch.no_grad()：PyTorch 將不再計(jì)算梯度，這將使得模型 forward 的時(shí)候，顯存的需求大幅減少，速度大幅提高。

注意：若模型中具有 Batch Normalization 操作，想固定該操作進(jìn)行訓(xùn)練時(shí)，需調(diào)用對應(yīng)的 module 的 eval() 函數(shù)。這是因?yàn)?BN Module 除了參數(shù)以外，還會(huì)對輸入的數(shù)據(jù)進(jìn)行統(tǒng)計(jì)，若不調(diào)用 eval()，統(tǒng)計(jì)量將發(fā)生改變！具體代碼可以這樣寫：

for module in model.modules():    module.eval()

在其他地方看到的解釋：

• model.eval() will notify all your layers that you are in eval mode, that way, batchnorm or dropout layers will work in eval model instead of training mode.

• torch.no_grad() impacts the autograd engine and deactivate it. It will reduce memory usage and speed up computations but you won’t be able to backprop (which you don’t want in an eval script).

6. 利用 torch.nn.DataParallel 進(jìn)行多 GPU 訓(xùn)練

import torchimport torch.nn as nnimport torchvision.models as models# 生成模型# 利用 torch.nn.DataParallel 進(jìn)行載入模型，默認(rèn)使用所有GPU（可以用 CUDA_VISIBLE_DEVICES 設(shè)置所使用的 GPU）model = nn.DataParallel(models.resnet18())
# 凍結(jié)參數(shù)for param in model.module.layer4.parameters():    param.requires_grad = Falseparam_optim = filter(lambda p:p.requires_grad, model.parameters())
# 設(shè)置測試模式model.module.layer4.eval()
# 保存模型參數(shù)（讀取所保存模型參數(shù)后，再進(jìn)行并行化操作，否則無法利用之前的代碼進(jìn)行讀取）torch.save(model.module.state_dict(),'./CheckPoint.pkl')

-完-