LR Range Test代码及画图

以下是一个示例的LR Range Test的代码，用于确定学习率的最佳范围：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

# 定义模型
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.linear = nn.Linear(1, 1)
        
    def forward(self, x):
        return self.linear(x)

# 定义数据集
class Dataset(torch.utils.data.Dataset):
    def __init__(self):
        self.x = torch.randn(100, 1)  # 输入特征
        self.y = 3 * self.x + 2 + torch.randn(100, 1)  # 输出标签
        
    def __getitem__(self, index):
        return self.x[index], self.y[index]
    
    def __len__(self):
        return len(self.x)

# LR Range Test函数
def lr_range_test(model, train_loader, optimizer, criterion, init_lr, final_lr, epochs):
    lr_values = []
    loss_values = []
    
    lr_scheduler = optim.lr_scheduler.LambdaLR(optimizer, lambda epoch: epoch)
    
    for epoch in range(epochs):
        for inputs, targets in train_loader:
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
            
            lr_values.append(lr_scheduler.get_last_lr()[0])
            loss_values.append(loss.item())
            
            lr_scheduler.step()
            
    return lr_values, loss_values

# 初始化模型、数据集和优化器
model = Model()
dataset = Dataset()
train_loader = DataLoader(dataset, batch_size=10, shuffle=True)
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=1e-6)

# 进行LR Range Test
lr_values, loss_values = lr_range_test(model, train_loader, optimizer, criterion, init_lr=1e-6, final_lr=10, epochs=100)

# 绘制LR Range Test结果
import matplotlib.pyplot as plt

plt.plot(lr_values, loss_values)
plt.xlabel('Learning Rate')
plt.ylabel('Loss')
plt.xscale('log')
plt.show()

这段代码演示了如何使用LR Range Test来确定学习率的最佳范围。在示例中，我们定义了一个简单的线性回归模型，并使用随机生成的数据集进行训练。LR Range Test会尝试一系列不同的学习率，然后记录每个学习率对应的损失值，最后绘制学习率和损失之间的关系图。

请注意，在实际使用中，你可能需要根据你的具体任务和模型进行一些调整和修改。此示例仅供参考。