def test_l1_norm_pruner(self):
model = TorchModel()
config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
pruner = L1NormPruner(model=model,
config_list=config_list,
mode='dependency_aware',
dummy_input=torch.rand(10, 1, 28, 28))
pruned_model, masks = pruner.compress()
pruner._unwrap_model()
sparsity_list = compute_sparsity_mask2compact(pruned_model, masks,
config_list)
assert 0.79 < sparsity_list[0]['total_sparsity'] < 0.81
def test_lottery_ticket_pruner(self):
model = TorchModel()
config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
pruner = LotteryTicketPruner(model,
config_list,
'level',
3,
log_dir='../../logs')
pruner.compress()
_, pruned_model, masks, _, _ = pruner.get_best_result()
sparsity_list = compute_sparsity_mask2compact(pruned_model, masks,
config_list)
assert 0.78 < sparsity_list[0]['total_sparsity'] < 0.82
def test_simulated_annealing_pruner(self):
model = TorchModel()
config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
pruner = SimulatedAnnealingPruner(model,
config_list,
evaluator,
start_temperature=40,
log_dir='../../logs')
pruner.compress()
_, pruned_model, masks, _, _ = pruner.get_best_result()
sparsity_list = compute_sparsity_mask2compact(pruned_model, masks,
config_list)
assert 0.78 < sparsity_list[0]['total_sparsity'] < 0.82
def test_admm_pruner(self):
model = TorchModel()
config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8, 'rho': 1e-3}]
pruner = ADMMPruner(model=model,
config_list=config_list,
trainer=trainer,
optimizer=get_optimizer(model),
criterion=criterion,
iterations=2,
training_epochs=1)
pruned_model, masks = pruner.compress()
pruner._unwrap_model()
sparsity_list = compute_sparsity_mask2compact(pruned_model, masks,
config_list)
assert 0.79 < sparsity_list[0]['total_sparsity'] < 0.81
def test_taylor_fo_pruner(self):
model = TorchModel()
config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
pruner = TaylorFOWeightPruner(model=model,
config_list=config_list,
trainer=trainer,
optimizer=get_optimizer(model),
criterion=criterion,
training_batches=1,
mode='dependency_aware',
dummy_input=torch.rand(10, 1, 28, 28))
pruned_model, masks = pruner.compress()
pruner._unwrap_model()
sparsity_list = compute_sparsity_mask2compact(pruned_model, masks,
config_list)
assert 0.79 < sparsity_list[0]['total_sparsity'] < 0.81
def test_slim_pruner(self):
model = TorchModel()
config_list = [{'op_types': ['BatchNorm2d'], 'total_sparsity': 0.8}]
pruner = SlimPruner(model=model,
config_list=config_list,
trainer=trainer,
optimizer=get_optimizer(model),
criterion=criterion,
training_epochs=1,
scale=0.001,
mode='global')
pruned_model, masks = pruner.compress()
pruner._unwrap_model()
sparsity_list = compute_sparsity_mask2compact(pruned_model, masks,
config_list)
assert 0.79 < sparsity_list[0]['total_sparsity'] < 0.81
def test_movement_pruner(self):
model = TorchModel()
config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
pruner = MovementPruner(model=model,
config_list=config_list,
trainer=trainer,
traced_optimizer=get_optimizer(model),
criterion=criterion,
training_epochs=5,
warm_up_step=0,
cool_down_beginning_step=4)
pruned_model, masks = pruner.compress()
pruner._unwrap_model()
sparsity_list = compute_sparsity_mask2compact(pruned_model, masks,
config_list)
assert 0.78 < sparsity_list[0]['total_sparsity'] < 0.82
def test_activation_mean_rank_pruner(self):
model = TorchModel()
config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
pruner = ActivationMeanRankPruner(
model=model,
config_list=config_list,
trainer=trainer,
traced_optimizer=get_optimizer(model),
criterion=criterion,
training_batches=5,
activation='relu',
mode='dependency_aware',
dummy_input=torch.rand(10, 1, 28, 28))
pruned_model, masks = pruner.compress()
pruner._unwrap_model()
sparsity_list = compute_sparsity_mask2compact(pruned_model, masks,
config_list)
assert 0.78 < sparsity_list[0]['total_sparsity'] < 0.82
def test_auto_compress_pruner(self):
model = TorchModel()
config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
admm_params = {
'trainer': trainer,
'optimizer': get_optimizer(model),
'criterion': criterion,
'iterations': 10,
'training_epochs': 1
}
sa_params = {
'evaluator': evaluator,
'start_temperature': 40
}
pruner = AutoCompressPruner(model, config_list, 10, admm_params, sa_params=sa_params, log_dir='../../../logs')
pruner.compress()
_, pruned_model, masks, _, _ = pruner.get_best_result()
sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
print(sparsity_list)
assert 0.78 < sparsity_list[0]['total_sparsity'] < 0.82
