在Pytorch中为不同层设置不同学习率来提升性能，优化深度学习模型

import torch
 import torch.nn as nn
 import torchvision.models as models
 
 # 加载预训练的ResNet18模型
 model = models.resnet18(pretrained=True)
 
 # 修改最后的全连接层以适应新的分类任务
 num_classes = 10  # 假设新任务有10个类别
 model.fc = nn.Linear(model.fc.in_features, num_classes)
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.

# 定义不同组的学习率
 backbone_lr = 1e-4  # 较小的学习率用于预训练的主干网络
 classifier_lr = 1e-3  # 较大的学习率用于新的分类器层
 
 # 创建参数组
 params = [
    {'params': model.conv1.parameters(), 'lr': backbone_lr},
    {'params': model.bn1.parameters(), 'lr': backbone_lr},
    {'params': model.layer1.parameters(), 'lr': backbone_lr},
    {'params': model.layer2.parameters(), 'lr': backbone_lr},
    {'params': model.layer3.parameters(), 'lr': backbone_lr},
    {'params': model.layer4.parameters(), 'lr': backbone_lr},
    {'params': model.fc.parameters(), 'lr': classifier_lr}
 ]
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.

optimizer = torch.optim.Adam(params)
1.

# 定义损失函数
 criterion = nn.CrossEntropyLoss()
 
 # 训练循环
 for epoch in range(num_epochs):
     model.train()
     for inputs, labels in train_loader:
         optimizer.zero_grad()
         outputs = model(inputs)
         loss = criterion(outputs, labels)
         loss.backward()
         optimizer.step()
     
     # 在每个epoch结束后进行验证
     model.eval()
     # ... [验证代码]
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.

from torch.optim.lr_scheduler import StepLR
 
 scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
 
 # 在训练循环中更新学习率
 for epoch in range(num_epochs):
     # ... [训练代码]
     scheduler.step()
1.
2.
3.
4.
5.
6.
7.
8.

# 初始阶段：只训练分类器
 for param in model.parameters():
     param.requires_grad = False
 model.fc.requires_grad = True
 
 # 训练几个epoch后
 model.layer4.requires_grad = True
 
 # 再过几个epoch
 model.layer3.requires_grad = True
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.

class LayerAdaptiveLR(torch.optim.Adam):
     def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
         super().__init__(params, lr, betas, eps, weight_decay)
         self.param_groups = sorted(self.param_groups, key=lambda x: id(x['params'][0]))
         
     def step(self, closure=None):
         loss = None
         if closure is not None:
             loss = closure()
 
         for group in self.param_groups:
             for p in group['params']:
                 if p.grad is None:
                     continue
                 grad = p.grad.data
                 state = self.state[p]
 
                 # 根据梯度统计调整学习率
                 if len(state) == 0:
                     state['step'] = 0
                     state['exp_avg'] = torch.zeros_like(p.data)
                     state['exp_avg_sq'] = torch.zeros_like(p.data)
 
                 exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
                 beta1, beta2 = group['betas']
 
                 state['step'] += 1
 
                 exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
                 exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
 
                 denom = exp_avg_sq.sqrt().add_(group['eps'])
                 
                 # 动态调整学习率
                 step_size = group['lr'] * (exp_avg.abs() / denom).mean().item()
                 p.data.add_(exp_avg, alpha=-step_size)
 
         return loss
 
 # 使用示例
 optimizer = LayerAdaptiveLR(model.parameters(), lr=1e-3)
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.