本文用於記錄如何進行 PyTorch 所提供的預訓練模型應如何加載，所訓練模型的參數應如何保存與讀取，如何凍結模型部分參數以方便進行 fine-tuning 以及如何利用多 GPU 訓練模型。
Update 2021.10.11 : 向大家推薦一個預訓練模型的論文庫，不僅可以查看相關的論文和全文，還能看到每篇論文的視頻解讀、研究溯源、代碼鏈接和自動提取的論文摘要等內容。
1. PyTorch 預訓練模型
Pytorch 提供了許多 Pre-Trained Model on ImageNet，僅需調用 torchvision.models 即可，具體細節可查看官方文檔。
往往我們需要對 Pre-Trained Model 進行相應的修改，以適應我們的任務。這種情況下，我們可以先輸出 Pre-Trained Model 的結構，確定好對哪些層修改，或者添加哪些層，接着，再將其修改即可。
比如，我需要將 ResNet-50 的 Layer 3 后的所有層去掉，在分別連接十個分類器，分類器由 ResNet-50.layer4 和 AvgPool Layer 和 FC Layer 構成。這里就需要用到 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)] 是一個python列表，必須要把它轉換成一個Module Llist列表才可以被 PyTorch 使用，否則在運行的時候會報錯: RuntimeError: Input type (CUDAFloatTensor) and weight type (CPUFloatTensor) should be the same
2. 保存模型參數
PyTorch 中保存模型的方式有許多種：
# 保存整個網絡 torch.save(model, PATH) # 保存網絡中的參數, 速度快，占空間少 torch.save(model.state_dict(),PATH) # 選擇保存網絡中的一部分參數或者額外保存其余的參數 torch.save({'state_dict': model.state_dict(), 'fc_dict':model.fc.state_dict(), 'optimizer': optimizer.state_dict(),'alpha': loss.alpha, 'gamma': loss.gamma}, PATH)3. 讀取模型參數
同樣的，PyTorch 中讀取模型參數的方式也有許多種：
# 讀取整個網絡 model = torch.load(PATH) # 讀取 Checkpoint 中的網絡參數 model.load_state_dict(torch.load(PATH)) # 若 Checkpoint 中的網絡參數與當前網絡參數有部分不同,有以下兩種方式進行加載: # 1. 利用字典的 update 方法進行加載 Checkpoint = torch.load(Path) model_dict = model.state_dict() model_dict.update(Checkpoint) model.load_state_dict(model_dict) # 2. 利用 load_state_dict() 的 strict 參數進行部分加載 model.load_state_dict(torch.load(PATH), strict=False)4. 凍結部分模型參數，進行 fine-tuning
加載完 Pre-Trained Model 后，我們需要對其進行 Finetune。但是在此之前，我們往往需要凍結一部分的模型參數：
# 第一種方式 for p in freeze.parameters(): # 將需要凍結的參數的 requires_grad 設置為 False p.requires_grad = False for p in no_freeze.parameters(): # 將fine-tuning 的參數的 requires_grad 設置為 True p.requires_grad = True # 將需要 fine-tuning 的參數放入optimizer 中 optimizer.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3) # 第二種方式 optim_param = [] for p in freeze.parameters(): # 將需要凍結的參數的 requires_grad 設置為 False p.requires_grad = False for p in no_freeze.parameters(): # 將fine-tuning 的參數的 requires_grad 設置為 True p.requires_grad = True optim_param.append(p) optimizer.SGD(optim_param, lr=1e-3) # 將需要 fine-tuning 的參數放入optimizer 中5. 模型訓練與測試的設置
訓練時，應調用 model.train() ；測試時，應調用 model.eval()，以及 with torch.no_grad():
model.train()：使 model 變成訓練模式，此時 dropout 和 batch normalization 的操作在訓練起到防止網絡過擬合的問題。
model.eval()：PyTorch會自動把 BN 和 DropOut 固定住，不會取平均，而是用訓練好的值。不然的話，一旦測試集的 Batch Size 過小，很容易就會被 BN 層導致生成圖片顏色失真極大。
with torch.no_grad()：PyTorch 將不再計算梯度，這將使得模型 forward 的時候，顯存的需求大幅減少，速度大幅提高。
注意：若模型中具有 Batch Normalization 操作，想固定該操作進行訓練時，需調用對應的 module 的 eval() 函數。這是因為 BN Module 除了參數以外，還會對輸入的數據進行統計，若不調用 eval()，統計量將發生改變！具體代碼可以這樣寫：
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 進行多 GPU 訓練
import torch import torch.nn as nn import torchvision.models as models # 生成模型 # 利用 torch.nn.DataParallel 進行載入模型，默認使用所有GPU（可以用 CUDA_VISIBLE_DEVICES 設置所使用的 GPU） model = nn.DataParallel(models.resnet18()) # 凍結參數 for param in model.module.layer4.parameters(): param.requires_grad = False param_optim = filter(lambda p:p.requires_grad, model.parameters()) # 設置測試模式 model.module.layer4.eval() # 保存模型參數（讀取所保存模型參數后，再進行並行化操作，否則無法利用之前的代碼進行讀取） torch.save(model.module.state_dict(),'./CheckPoint.pkl')