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
