LR Range Test代码

当进行模型的训练时，我们常常需要进行学习率范围测试（LR Range Test）来选择合适的学习率。下面是一个示例的LR Range Test代码，使用PyTorch框架：

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

# 定义你的模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc = nn.Linear(10, 1)

    def forward(self, x):
        return self.fc(x)

# 定义你的数据集
class MyDataset(torch.utils.data.Dataset):
    def __init__(self):
        self.data = torch.randn((100, 10))
        self.labels = torch.randn((100, 1))

    def __getitem__(self, index):
        return self.data[index], self.labels[index]

    def __len__(self):
        return len(self.data)

# 定义LR Range Test函数
def lr_range_test(model, train_loader, optimizer, criterion, init_value=1e-8, final_value=10.0, beta=0.98):
    num_batches = len(train_loader)
    lr_lambda = lambda iteration: (final_value / init_value) ** (1 / num_batches)
    lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

    model.train()
    avg_loss = 0.0
    best_loss = float('inf')
    smooth_loss = 0.0

    for batch_index, (inputs, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        
        avg_loss = beta * avg_loss + (1 - beta) * loss.item()
        smooth_loss = avg_loss / (1 - beta ** (batch_index + 1))

        if smooth_loss < best_loss:
            best_loss = smooth_loss

        loss.backward()
        optimizer.step()

        lr_scheduler.step()

    return best_loss

# 准备数据
dataset = MyDataset()
train_loader = DataLoader(dataset, batch_size=10, shuffle=True)

# 初始化模型、优化器和损失函数
model = MyModel()
optimizer = optim.SGD(model.parameters(), lr=1e-4)
criterion = nn.MSELoss()

# 进行LR Range Test
best_loss = lr_range_test(model, train_loader, optimizer, criterion)

print("Best loss:", best_loss)

请注意，这只是一个简单的示例，具体的LR Range Test实现可能因任务的不同而有所变化。你可以根据自己的需求进行相应的修改和调整。