def test_suggest_lr(): task = mod_task.XORTask() lr_finder = prepare_lr_finder(task) lr_finder.history["loss"] = [10, 8, 4, 1, 4, 16] lr_finder.history["lr"] = range(len(lr_finder.history["loss"])) fig, ax = plt.subplots() ax, lr = lr_finder.plot(skip_start=0, skip_end=0, suggest_lr=True, ax=ax) assert lr == 2 # Loss with minimal gradient is the first element in history lr_finder.history["loss"] = [1, 0, 1, 2, 3, 4] lr_finder.history["lr"] = range(len(lr_finder.history["loss"])) fig, ax = plt.subplots() ax, lr = lr_finder.plot(skip_start=0, skip_end=0, suggest_lr=True, ax=ax) assert lr == 0 # Loss with minimal gradient is the last element in history lr_finder.history["loss"] = [0, 1, 2, 3, 4, 3] lr_finder.history["lr"] = range(len(lr_finder.history["loss"])) fig, ax = plt.subplots() ax, lr = lr_finder.plot(skip_start=0, skip_end=0, suggest_lr=True, ax=ax) assert lr == len(lr_finder.history["loss"]) - 1
def test_plot_with_skip_and_suggest_lr(suggest_lr, skip_start, skip_end): task = mod_task.XORTask() num_iter = 11 # prepare_lr_finder sets the starting lr to 1e-5 lr_finder = prepare_lr_finder(task) lr_finder.range_test(task.train_loader, num_iter=num_iter, step_mode="exp", end_lr=0.1) fig, ax = plt.subplots() results = lr_finder.plot(skip_start=skip_start, skip_end=skip_end, suggest_lr=suggest_lr, ax=ax) if num_iter - skip_start - skip_end <= 1: # handle data with one or zero lr assert len(ax.lines) == 1 assert results is ax else: # handle different suggest_lr # for 'steepest': the point with steepest gradient (minimal gradient) assert len(ax.lines) == 1 assert len(ax.collections) == int(suggest_lr) if results is not ax: assert len(results) == 2
def test_scheduler_and_num_iter(num_iter, scheduler): task = mod_task.XORTask() # prepare_lr_finder sets the starting lr to 1e-5 lr_finder = prepare_lr_finder(task) with pytest.raises(ValueError, match="num_iter"): lr_finder.range_test(task.train_loader, num_iter=num_iter, step_mode=scheduler, end_lr=5e-5)
def test_run_range_test_with_traindataloaderiter(self, mocker): task = mod_task.XORTask() lr_finder = prepare_lr_finder(task) num_iter = 5 loader_iter = TrainDataLoaderIter(task.train_loader) spy = mocker.spy(loader_iter, "inputs_labels_from_batch") lr_finder.range_test(loader_iter, num_iter=num_iter) assert spy.call_count == num_iter
def test_run_range_test_with_trainloaderiter_without_subclassing(self): task = mod_task.XORTask() lr_finder = prepare_lr_finder(task) num_iter = 5 loader_iter = CustomLoaderIter(task.train_loader) with pytest.raises(ValueError, match="`train_loader` has unsupported type"): lr_finder.range_test(loader_iter, num_iter=num_iter)
def test_suggest_lr(): task = mod_task.XORTask() lr_finder = prepare_lr_finder(task) lr_finder.history["loss"] = [10, 8, 4, 1, 4, 16] lr_finder.history["lr"] = range(len(lr_finder.history["loss"])) fig, ax = plt.subplots() ax, lr = lr_finder.plot(skip_start=0, skip_end=0, suggest_lr=True, ax=ax) assert lr == 2
def test_exponential_lr_history(self): task = mod_task.XORTask() # prepare_lr_finder sets the starting lr to 1e-5 lr_finder = prepare_lr_finder(task) lr_finder.range_test(task.train_loader, num_iter=5, step_mode="exp", end_lr=0.1) assert len(lr_finder.history["lr"]) == 5 assert lr_finder.history["lr"] == pytest.approx( [1e-5, 1e-4, 1e-3, 1e-2, 0.1])
def test_linear_lr_history(self): task = mod_task.XORTask() # prepare_lr_finder sets the starting lr to 1e-5 lr_finder = prepare_lr_finder(task) lr_finder.range_test(task.train_loader, num_iter=5, step_mode="linear", end_lr=5e-5) assert len(lr_finder.history["lr"]) == 5 assert lr_finder.history["lr"] == pytest.approx( [1e-5, 2e-5, 3e-5, 4e-5, 5e-5])
def test_run_range_test_with_valloaderiter_without_subclassing(self): task = mod_task.XORTask(validate=True) lr_finder = prepare_lr_finder(task) num_iter = 5 train_loader_iter = TrainDataLoaderIter(task.train_loader) val_loader_iter = CustomLoaderIter(task.val_loader) with pytest.raises(ValueError, match="`val_loader` has unsupported type"): lr_finder.range_test(train_loader_iter, val_loader=val_loader_iter, num_iter=num_iter)
def test_run_range_test_with_valdataloaderiter(self, mocker): task = mod_task.XORTask(validate=True) lr_finder = prepare_lr_finder(task) num_iter = 5 train_loader_iter = TrainDataLoaderIter(task.train_loader) val_loader_iter = ValDataLoaderIter(task.val_loader) spy_train = mocker.spy(train_loader_iter, "inputs_labels_from_batch") spy_val = mocker.spy(val_loader_iter, "inputs_labels_from_batch") lr_finder.range_test(train_loader_iter, val_loader=val_loader_iter, num_iter=num_iter) assert spy_train.call_count == num_iter assert spy_val.call_count == num_iter * len(task.val_loader)
def test_gradient_accumulation(self, mocker): desired_bs, accum_steps = 32, 4 real_bs = desired_bs // accum_steps num_iter = 10 task = mod_task.XORTask(batch_size=real_bs) lr_finder = prepare_lr_finder(task) spy = mocker.spy(lr_finder, "criterion") lr_finder.range_test(task.train_loader, num_iter=num_iter, accumulation_steps=accum_steps) # NOTE: We are using smaller batch size to simulate a large batch. # So that the actual times of model/criterion called should be # `(desired_bs/real_bs) * num_iter` == `accum_steps * num_iter` assert spy.call_count == accum_steps * num_iter
def test_mixed_precision(self, mocker): batch_size = 32 num_iter = 10 task = mod_task.XORTask(batch_size=batch_size) # Wrap model and optimizer by `amp.initialize`. Beside, `amp` requires # CUDA GPU. So we have to move model to GPU first. model, optimizer, device = task.model, task.optimizer, task.device model = model.to(device) task.model, task.optimizer = amp.initialize(model, optimizer) assert hasattr(task.optimizer, "_amp_stash") lr_finder = prepare_lr_finder(task) spy = mocker.spy(amp, "scale_loss") lr_finder.range_test(task.train_loader, num_iter=num_iter) # NOTE: Here we did not perform gradient accumulation, so that call count # of `amp.scale_loss` should equal to `num_iter`. assert spy.call_count == num_iter
def test_gradient_accumulation_with_apex_amp(self, mocker): desired_bs, accum_steps = 32, 4 real_bs = desired_bs // accum_steps num_iter = 10 task = mod_task.XORTask(batch_size=real_bs) # Wrap model and optimizer by `amp.initialize`. Beside, `amp` requires # CUDA GPU. So we have to move model to GPU first. model, optimizer, device = task.model, task.optimizer, task.device model = model.to(device) task.model, task.optimizer = amp.initialize(model, optimizer) lr_finder = prepare_lr_finder(task) spy = mocker.spy(amp, "scale_loss") lr_finder.range_test(task.train_loader, num_iter=num_iter, accumulation_steps=accum_steps) assert spy.call_count == accum_steps * num_iter