Scheduler'ы в pytorch слишком сильно занижают lr?

Question

[ermak148]

Переучиваюсь с tensorflow на torch. Вспомнил про какие-то scheduler'ы при обучении gpt-2. Написал функцию и засунул в LambdaLR. Отслеживаю lr во время обучения и вижу что он очень сильно упал (0.0000005) при определенном минимуме в самой функции. Такое же происходит с StepLR ( другие не пробовал ) (Тыкал на оптимизаторах AdamW и RMSprop) (функция с обычным принтом в тесте отлично работала!).
Сама функция и остальной код:

def get_lr(epoch):
    min_lr = 0.00001
    max_lr = 0.001
    step = 0.00002

    return max(min_lr, max_lr - step * epoch)
optim = torch.optim.RMSprop(model.parameters(), lr=0.001)
# scheduler = LambdaLR(optim, get_lr)
with tqdm.tqdm(total=epochs, postfix={"Loss": 0, "LR": 0.0009}, dynamic_ncols=True) as pbar:
    for epoch in range(epochs):
        for inputs, targets in dataloader:
            # print(inputs, targets)
            optim.zero_grad()
            out = model(inputs)
            loss = criterion(out, targets)
            loss.backward()
            optim.step()
        pbar.update(1)
        pbar.set_postfix({"Loss": str(loss.item()), "LR": "{:.10f}".format(optim.param_groups[0]['lr'])})
        # print(f"Epoch: {epoch}. Loss: {loss.item()}")
        # scheduler.step()

Хелп!

Answer 1

[Lapita12]

def get_lr(epoch):
    min_lr = 0.00001
    max_lr = 0.001
    step = 0.00002

    return max(min_lr, max_lr * (1 - step * epoch)**2)