1. 程式人生 > 程式設計 >pytorch 實現模型不同層設定不同的學習率方式

pytorch 實現模型不同層設定不同的學習率方式

在目標檢測的模型訓練中, 我們通常都會有一個特徵提取網路backbone, 例如YOLO使用的darknet SSD使用的VGG-16。

為了達到比較好的訓練效果, 往往會載入預訓練的backbone模型引數, 然後在此基礎上訓練檢測網路, 並對backbone進行微調, 這時候就需要為backbone設定一個較小的lr。

class net(torch.nn.Module):
  def __init__(self):
    super(net,self).__init__()
    # backbone
    self.backbone = ...
    # detect
    self....

在設定optimizer時, 只需要引數分為兩個部分, 並分別給定不同的學習率lr。

base_params = list(map(id,net.backbone.parameters()))
logits_params = filter(lambda p: id(p) not in base_params,net.parameters())
params = [
  {"params": logits_params,"lr": config.lr},{"params": net.backbone.parameters(),"lr": config.backbone_lr},]
optimizer = torch.optim.SGD(params,momentum=config.momentum,weight_decay=config.weight_decay)
 

以上這篇pytorch 實現模型不同層設定不同的學習率方式就是小編分享給大家的全部內容了,希望能給大家一個參考,也希望大家多多支援我們。