pytorch--训练分层学习率设置

在训练模型时,我们经常会使用两个神经网络模型进行融合,若两个模型的复杂度不同,或者激活函数不同,导致训练后的模型训练损失忽高忽低,差距巨大,有可能是陷入了局部最优的状况。这时候采用分层学习率的策略可能帮助模型度过局部最优困境。

下面是一个简单的示例:

对于一个继承于nn.Module的神经网络模型Model

class Model(nn.Module):
	def __init__(self):
		super().__init__()
		self.layer1 = nn.Sequential(nn.Linear(20, 10), nn.Tanh())
		self.layer2 = nn.Linear(10, 1)
	
	def forward(self, x):
		out = self.layer1(x)
		out = self.layer2(out)
		return out

那么分层学习率的设置大致如下:

model = Model() # 模型初始化
# 设置分层学习率
params_list = [{'params': model.layer1.parameters(), 'lr': 0.001},
	{'params': model.layer2.parameters(), 'lr': 0.002}]
# 将学习率传入优化器 
optimizeer = torch.optim.RMSprop(params_list)
# 模型训练
train(model, max_epoch, optimizer, train_iter, vali_iter, test_iter, loss_func)
posted @ 2023-05-27 21:44  woxin_lab  阅读(283)  评论(0编辑  收藏  举报