天天看點

基于pytorch的深度模型參數分别優化與當機

問題:

在深度學習的模型的訓練的過程中,經常會遇到需要對不同的模型參數設定不同的學習政策,比如使用不同的學習率,設定權重衰減等,在pytorch架構中的設定相對來說是比較簡單的。

pytorch中分别優化設定:

#coding:utf-8

import torch
from torchvision import models

#模型初始化
model = models.resnet50(pretrained=True)

#根據需要,進行參數分組(權重衰減設定)
pg0, pg1, pg2 = [], [], []  # optimizer parameter groups
for k, v in model.named_modules():
    if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
        pg2.append(v.bias)  # biases
    if isinstance(v, nn.BatchNorm2d):
        pg0.append(v.weight)  # no decay
    elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
        pg1.append(v.weight)  # apply decay

#優化器設定
optimizer = optim.SGD(pg0, lr=0.001, momentum=0.0005, nesterov=True)
optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']}) 
optimizer.add_param_group({'params': pg2})  # add pg2 (biases)

           

pytorch中當機某些層設定:

from torchvision import models

model = models.resnet18()

# 當機layer4以前的所有層,隻訓練layer4與全連接配接層(根據類别需自定義)
no_freeze = ['layer4','fc']  # parameter names to freeze (full or partial)

for k, v in model.named_parameters():
    v.requires_grad = True  # train all layers
    if not any(x in k for x in no_freeze):
        print('freezing %s' % k)
        v.requires_grad = False
           

參考:yolov5源碼庫中的train.py檔案中

繼續閱讀