Open in app

Sign in

Write

Sign in

Pytorch 凍結與解凍模型參數、layer 設置不同 learning rate

李謦伊
謦伊的閱讀筆記
Published in
15 min readAug 19, 2022

在進行 Transfer learning 時經常會使用 pre-trained model weight 做 feature extraction (特徵提取) 或 fine-tuning (微調)，feature extraction 的做法是將 輸出層之前的網路層凍結，在訓練自己的資料集時，只訓練更新後面未被凍結的網路層；而 fine-tuning 的做法則是訓練更新 pre-trained model 中的每一層 weight 或是僅凍結較低層的網路層參數，因此使用 fine-tuning 的訓練時間會比 feature extraction 長。

source: feature extraction; fine-tuning

本文將要來介紹用 pytorch 來進行凍結、解凍 pre-trained model 參數的方法以及如何在不同的網路層中設置不同的 learning rate，所有 code 會放在文章最下方。

首先 import 需要的 library

import torch
import torch.nn as nn
import torchvision.models as models

凍結 pre-trained model weight

可使用 pytorch 的 torchvision.models 中所提供的模型權重，也可以使用自己訓練或下載的模型權重檔。

🔖 使用 pytorch 提供的 pre-trained model 權重

  • 讀取模型
model = models.resnet18(pretrained=True)

查看模型網路層名稱和參數，其中 requires_grad 表示是否需要在計算中保留該參數的梯度。

for name, param in model_1.named_parameters():
    print(“name: “, name)
    print(“requires_grad: “, param.requires_grad)# === output ===
name:  conv1.weight
requires_grad:  True
name:  bn1.weight
requires_grad:  True
name:  bn1.bias
requires_grad:  True
name:  layer1.0.conv1.weight
requires_grad:  True
name:  layer1.0.bn1.weight
requires_grad:  True
name: fc.weight
requires_grad:  True
name: fc.bias
requires_grad:  True
  • 將除了全連接層以外的網路層都凍結
for name, param in model_1.named_parameters():
    if name not in [‘fc.weight’, ‘fc.bias’]:
    param.requires_grad = False

由下列結果可看到除了全連接層以外的參數的 requires_grad 變為 False

for name, param in model_1.named_parameters():
    print(“name: “, name)
    print(“requires_grad: “, param.requires_grad)# === output ===
name:  conv1.weight
requires_grad:  False
name:  bn1.weight
requires_grad:  False
name:  bn1.bias
requires_grad:  False
name:  layer1.0.conv1.weight
requires_grad:  False
name:  layer1.0.bn1.weight
requires_grad:  False
name: fc.weight
requires_grad:  True
name: fc.bias
requires_grad:  True
  • 在優化器中加入 filter 進行過濾

接著要將 requires_grad 為 False 的參數過濾掉，只傳入 requires_grad 為 True 的參數給 optimizer。

filter() 函數用於過濾不符合條件的元素，其語法為 filter(function, iterable)，其中 function 為判斷條件、iterable 為元素列表，返回值為一個迭代器。進行過濾的過程是將 iterable 的值傳遞給 function 進行判斷，返回值為 True 的元素才會放入要回傳的迭代器中。

來看個簡單的 filter() 函數範例，首先建立一個 function 為回傳奇數的數字，iterable 為 1~10 的 list，經過 filter() 後所得到的值為所有奇數的 list。

### filter exampledef is_odd(n):
    return n % 2 == 1
 
newlist = filter(is_odd, [i for i in range(10)])
print(list(newlist))# === output ===
[1, 3, 5, 7, 9]

將參數進行過濾後傳給 optimizer

parameters_1 = filter(lambda p: p.requires_grad, model_1.parameters())optimizer_1 = torch.optim.Adam(parameters_1, lr=0.001, weight_decay=1e-5)

🔖 使用自己訓練或下載的權重檔

以 resnet18 為例，首先下載 resnet18 的權重檔

import gdownresnet_model = ‘https://download.pytorch.org/models/resnet18-5c106cde.pth'
gdown.download(resnet_model, “resnet-5c106cde.pth”)
  • 讀取權重檔
checkpoint = torch.load(‘resnet-5c106cde.pth’)

查看模型網路層名稱和參數，其中 requires_grad 表示是否需要在計算中保留該參數的梯度。

for k, v in checkpoint.items():
    print(“name: “, k)
    print(“requires_grad: “, v.requires_grad)# === output ===
name:  conv1.weight
requires_grad:  True
name:  bn1.running_mean
requires_grad:  False
name:  bn1.running_var
requires_grad:  False
name:  bn1.weight
requires_grad:  True
name:  bn1.bias
requires_grad:  True
name: fc.weight
requires_grad:  True
name: fc.bias
requires_grad:  True
  • 將除了全連接層以外的網路層都凍結
for k, v in checkpoint.items():
    if k not in ['fc.weight', 'fc.bias']:
        v.requires_grad = False

由下列結果可看到除了全連接層以外的參數的 requires_grad 變為 False

for k, v in checkpoint.items():
    print("name: ", k)
    print("requires_grad: ", v.requires_grad)# === output ===
name:  conv1.weight
requires_grad:  False
name:  bn1.running_mean
requires_grad:  False
name:  bn1.running_var
requires_grad:  False
name:  bn1.weight
requires_grad:  False
name:  bn1.bias
requires_grad:  False
name: fc.weight
requires_grad:  True
name: fc.bias
requires_grad:  True
  • 在優化器中加入 filter 進行過濾

將 pre-trained model 中與模型相同網路層的權重參數提取至模型中，接著載入至模型中

model_2 = models.resnet18()
model_state = model_2.state_dict()pretrained_dict = {k: v for k, v in checkpoint.items() if k in model_state}
model_state.update(pretrained_dict)
model_2.load_state_dict(model_state)

將參數進行過濾後傳給 optimizer

parameters_2 = filter(lambda p: p.requires_grad, model_2.parameters())optimizer_2 = torch.optim.Adam(parameters_2, lr=0.001, weight_decay=1e-5)
  • 將 50 層內的 layer 凍結

若是想凍結 50 層內的網路層，可以採用以下方式

# 讀取模型
checkpoint = torch.load(‘resnet-5c106cde.pth’)# 凍結網路層
for i, (k, v) in enumerate(checkpoint.items()):
    if i < 50:
       v.requires_grad = False# 提取權重參數並載入模型
model_3 = models.resnet18()
model_state = model_3.state_dict()pretrained_dict = {k: v for k, v in checkpoint.items() if k in model_state}
model_state.update(pretrained_dict)
model_3.load_state_dict(model_state)# 將參數進行過濾後傳給 optimizer
parameters_3 = filter(lambda p: p.requires_grad, model_3.parameters())optimizer_3 = torch.optim.Adam(parameters_3, lr=0.001, weight_decay=1e-5)

解除凍結

很多時候並不知道該選取多少層的網路層進行凍結，可使用漸進式微調的方法來找出適合的微調層數。做法為先凍結某些部分的網路層，在進行一次 epoch 後再一步步地一層一層解凍網路層，並同時計算 loss，直到 loss 穩定以決定最終微調層數。

查看網路層的 requires_grad 狀態

for name, param in model_1.named_parameters():
    print("name: ", name)
    print("requires_grad: ", param.requires_grad)# === output ===
name:  conv1.weight
requires_grad:  False
name:  bn1.weight
requires_grad:  False
name:  bn1.bias
requires_grad:  False
name:  layer1.0.conv1.weight
requires_grad:  False
name:  layer1.0.bn1.weight
requires_grad:  False
  • 解除凍結
for name, param in model_1.named_parameters():
    if param.requires_grad == False:
        param.requires_grad = True
        optimizer_1.add_param_group({'params': param})

由下列結果可看到網路層參數的 requires_grad 變為 True

for name, param in model_1.named_parameters():
    print("name: ", name)
    print("requires_grad: ", param.requires_grad)# === output ===
name:  conv1.weight
requires_grad:  True
name:  bn1.weight
requires_grad:  True
name:  bn1.bias
requires_grad:  True
name:  layer1.0.conv1.weight
requires_grad:  True
name:  layer1.0.bn1.weight
requires_grad:  True

在不同的網路層中設置不同的 learning rate

若需要針對不同的網路層來設定各自的 learning rate，可採用以下三種方法:

  • 直接設定參數名稱
  • 設定某一層的全部參數
  • 將所有 layer 拆分成幾部分後再設定

首先查看一下網路層名稱

for name, param in model_1.named_parameters():
    print(name)# === output ===
conv1.weight
bn1.weight
bn1.bias
layer1.0.conv1.weight
layer1.0.bn1.weight
layer1.0.bn1.bias
...
layer4.1.conv2.weight
layer4.1.bn2.weight
layer4.1.bn2.bias
fc.weight
fc.bias

直接設定參數名稱

能夠對於全連接層中的 weight、bias 來設定不同的 learning rate，首先把要設定的參數寫成 {‘params’: <model_weight>, ‘lr’: <lr_value>} ，再將多個設定值包成 list 傳入 optimizer，另外第四行的 learning rate 為默認的值。

optimizer_4 = torch.optim.Adam([
    {'params': model_1.fc.weight, 'lr': 1e-5},
    {'params': model_1.fc.bias, 'lr': 1e-3}],
    lr=0.001) # 默認 learning rate 值

設定某一層的全部參數

也可直接設定某一層的全部參數，寫法如下所示

optimizer_5 = torch.optim.Adam([
    {'params': model_1.layer1.parameters(), 'lr': 1e-5},
    {'params': model_1.fc.parameters(), 'lr': 1e-3}],
    lr=0.001) # 默認 learning rate 值

將所有 layer 拆分成幾部分後再設定

假設要將網路層分成兩個部分：最後一層的全連接層和其他 CNN 層

列出所有全連接層的參數 id，然後用 filter() 過濾掉就能得到除了全連接層外的所有參數。

fc_layer = list(map(id, model_1.fc.parameters()))other_layer = filter(lambda p: id(p) not in fc_layer, model_1.parameters())

全連接層的參數與上方的相同為 model_1.fc.parameters()，而其他 CNN 層則是剛利用 filter() 過濾掉所得到的值。

optimizer_6 = torch.optim.Adam([
    {'params': other_layer, 'lr': 1e-5},
    {'params': model_1.fc.parameters(), 'lr': 1e-3}],
    lr=0.001) # 默認 learning rate 值

完整的 Code 可參考我的 Github:

Pytorch_note/freeze_unfreeze_network.ipynb at master · chingi071/Pytorch_note

Contribute to chingi071/Pytorch_note development by creating an account on GitHub.

github.com

Reference

Pytorch 相關文章

若我的文章對你有幫助的話，可以按下方的 “Give a tip” 按鈕給予我支持與鼓勵，謝謝 :)

Pytorch
Deep Learning
Artificial Intelligence

--

--

李謦伊
謦伊的閱讀筆記

Written by 李謦伊

4K Followers
Editor for

AI Engineer, Enjoy Deep Learning, Machine Learning～Contact me: chingi.lee071@gmail.com

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams