def test_concat_scheduler_state_dict(): tensor = torch.zeros([1], requires_grad=True) optimizer = torch.optim.SGD([tensor], lr=0) scheduler_1 = LinearCyclicalScheduler(optimizer, "lr", start_value=1.0, end_value=0.0, cycle_size=10) scheduler_2 = CosineAnnealingScheduler(optimizer, "lr", start_value=0.0, end_value=1.0, cycle_size=10) durations = [10] concat_scheduler = ConcatScheduler(schedulers=[scheduler_1, scheduler_2], durations=durations, save_history=False) state_dict = concat_scheduler.state_dict() assert state_dict["durations"] == durations assert state_dict["_current_duration"] == durations[0] assert state_dict["_scheduler_index"] == 0 for _ in range(20): concat_scheduler(None, None) concat_scheduler.load_state_dict(state_dict) assert concat_scheduler.durations == durations assert concat_scheduler._current_duration == durations[0] assert id(concat_scheduler._current_scheduler) == id(scheduler_1) with pytest.raises(ValueError, match=r"Required state attribute 'schedulers' is absent in provided state_dict"): concat_scheduler.load_state_dict({"a": 1}) with pytest.raises(ValueError, match=r"Input state_dict contains 0 state_dicts of concatenated schedulers"): concat_scheduler.load_state_dict({"schedulers": []}) with pytest.raises(TypeError, match=r"Argument state_dict should be a dictionary, but given"): concat_scheduler.load_state_dict(None)
def test_concat_scheduler_3_schedulers(): tensor = torch.zeros([1], requires_grad=True) optimizer = torch.optim.SGD([tensor], lr=0) scheduler_1 = LinearCyclicalScheduler(optimizer, "lr", start_value=1.0, end_value=0.5, cycle_size=20) scheduler_2 = LinearCyclicalScheduler(optimizer, "lr", start_value=0.5, end_value=0.45, cycle_size=10) scheduler_3 = LinearCyclicalScheduler(optimizer, "lr", start_value=0.5, end_value=0.0, cycle_size=20) durations = [10, 5] concat_scheduler = ConcatScheduler( schedulers=[scheduler_1, scheduler_2, scheduler_3], durations=durations, save_history=True ) state_dict = concat_scheduler.state_dict() data = [0] * 10 max_epochs = 2 simulated_values = ConcatScheduler.simulate_values( num_events=len(data) * max_epochs, schedulers=[scheduler_1, scheduler_2, scheduler_3], durations=durations ) def save_lr(engine): lrs.append(optimizer.param_groups[0]["lr"]) trainer = Engine(lambda engine, batch: None) trainer.add_event_handler(Events.ITERATION_STARTED, concat_scheduler) trainer.add_event_handler(Events.ITERATION_COMPLETED, save_lr) for _ in range(2): lrs = [] trainer.run(data, max_epochs=max_epochs) assert lrs == list( map( pytest.approx, [ # Cycle 1 of the first LinearCyclicalScheduler 1.0, 0.95, 0.9, 0.85, 0.8, 0.75, 0.7, 0.65, 0.6, 0.55, # Cycle 1 of the second LinearCyclicalScheduler 0.5, 0.49, 0.48, 0.47, 0.46, # Cycle 1 of the third LinearCyclicalScheduler 0.5, 0.45, 0.4, 0.35, 0.3, ], ) ) state_lrs = trainer.state.param_history["lr"] assert len(state_lrs) == len(lrs) # Unpack singleton lists assert [group[0] for group in state_lrs] == lrs assert lrs == pytest.approx([v for i, v in simulated_values]) concat_scheduler.load_state_dict(state_dict) trainer.state.param_history = None
def test_concat_scheduler_two_schedulers(duration_vals_as_np_int): tensor = torch.zeros([1], requires_grad=True) optimizer = torch.optim.SGD([tensor], lr=0) scheduler_1 = LinearCyclicalScheduler(optimizer, "lr", start_value=1.0, end_value=0.0, cycle_size=10) scheduler_2 = CosineAnnealingScheduler(optimizer, "lr", start_value=0.0, end_value=1.0, cycle_size=10) durations = [10] if duration_vals_as_np_int: durations = [np.int64(t) for t in durations] concat_scheduler = ConcatScheduler(schedulers=[scheduler_1, scheduler_2], durations=durations, save_history=True) state_dict = concat_scheduler.state_dict() data = [0] * 10 max_epochs = 2 simulated_values = ConcatScheduler.simulate_values( num_events=len(data) * max_epochs, schedulers=[scheduler_1, scheduler_2], durations=durations ) def save_lr(engine): lrs.append(optimizer.param_groups[0]["lr"]) trainer = Engine(lambda engine, batch: None) trainer.add_event_handler(Events.ITERATION_STARTED, concat_scheduler) trainer.add_event_handler(Events.ITERATION_COMPLETED, save_lr) for _ in range(2): lrs = [] trainer.run(data, max_epochs=max_epochs) assert lrs == list( map( pytest.approx, [ # Cycle 1 of the LinearCyclicalScheduler 1.0, 0.8, 0.6, 0.4, 0.2, 0.0, 0.2, 0.4, 0.6, 0.8, # Cycle 1 of the CosineAnnealingScheduler 0.0, 0.02447174185242318, 0.09549150281252627, 0.20610737385376332, 0.3454915028125263, 0.5, 0.6545084971874737, 0.7938926261462365, 0.9045084971874737, 0.9755282581475768, ], ) ) state_lrs = trainer.state.param_history["lr"] assert len(state_lrs) == len(lrs) # Unpack singleton lists assert [group[0] for group in state_lrs] == lrs assert lrs == pytest.approx([v for i, v in simulated_values]) concat_scheduler.load_state_dict(state_dict) trainer.state.param_history = None