Function test_lr_scheduler

tests/test_lr_scheduler.py:9–27

()

def test_lr_scheduler():
    warmup_steps = 200
    model = resnet50().cuda()
    optimizer = Adam(model.parameters(), lr=0.01)
    scheduler = LinearWarmupLR(optimizer, warmup_steps=warmup_steps)
    current_lr = scheduler.get_lr()[0]
    data = torch.rand(1, 3, 224, 224).cuda()

    for i in tqdm(range(warmup_steps * 2)):
        out = model(data)
        out.mean().backward()
        optimizer.step()
        scheduler.step()

        if i >= warmup_steps:
            assert scheduler.get_lr()[0] == 0.01
        else:
            assert scheduler.get_lr()[0] > current_lr, f"{scheduler.get_lr()[0]} <= {current_lr}"
            current_lr = scheduler.get_lr()[0]


if __name__ == "__main__":