def test_pruning_str_structured(): dataloader = prepare_experiment() model = nn.Linear(100, 10, bias=False) runner = dl.SupervisedRunner() criterion = nn.CrossEntropyLoss() runner.train( model=model, optimizer=torch.optim.Adam(model.parameters()), criterion=criterion, loaders={"train": dataloader}, callbacks=[PruningCallback("ln_structured", dim=1, l_norm=2)], num_epochs=1, ) assert np.isclose(pruning_factor(model), 0.5)
def test_parametrization(): dataloader = prepare_experiment() model = nn.Linear(100, 10, bias=False) runner = dl.SupervisedRunner() criterion = nn.CrossEntropyLoss() runner.train( model=model, optimizer=torch.optim.Adam(model.parameters()), criterion=criterion, loaders={"train": dataloader}, callbacks=[ PruningCallback(l1_unstructured, remove_reparametrization_on_stage_end=False) ], num_epochs=1, ) assert np.isclose(pruning_factor(model), 0.5) try: _mask = model.weight_mask mask_applied = True except AttributeError: mask_applied = False assert mask_applied