def test_channel_prune(self):
orig_net = resnet18(num_classes=10).to(device)
channel_prune(orig_net)
state_dict = torch.load(MODEL_FILE)
orig_net = resnet18(num_classes=10).to(device)
orig_net.load_state_dict(state_dict)
apply_compression_results(orig_net, MASK_FILE)
orig_net.eval()
net = resnet18(num_classes=10).to(device)
net.load_state_dict(state_dict)
net.eval()
data = torch.randn(BATCH_SIZE, 3, 224, 224).to(device)
ms = ModelSpeedup(net, data, MASK_FILE)
ms.speedup_model()
ms.bound_model(data)
net.eval()
ori_sum = orig_net(data).abs().sum().item()
speeded_sum = net(data).abs().sum().item()
print(ori_sum, speeded_sum)
assert (abs(ori_sum - speeded_sum) / abs(ori_sum) < RELATIVE_THRESHOLD) or \
(abs(ori_sum - speeded_sum) < ABSOLUTE_THRESHOLD)
def test_speedup_bigmodel(self):
prune_model_l1(BigModel())
model = BigModel()
apply_compression_results(model, MASK_FILE, 'cpu')
model.eval()
mask_out = model(dummy_input)
model.train()
ms = ModelSpeedup(model, dummy_input, MASK_FILE)
ms.speedup_model()
assert model.training
model.eval()
speedup_out = model(dummy_input)
if not torch.allclose(mask_out, speedup_out, atol=1e-07):
print('input:', dummy_input.size(),
torch.abs(dummy_input).sum((2, 3)))
print('mask_out:', mask_out)
print('speedup_out:', speedup_out)
raise RuntimeError('model speedup inference result is incorrect!')
orig_model = BigModel()
assert model.backbone2.conv1.out_channels == int(
orig_model.backbone2.conv1.out_channels * SPARSITY)
assert model.backbone2.conv2.in_channels == int(
orig_model.backbone2.conv2.in_channels * SPARSITY)
assert model.backbone2.conv2.out_channels == int(
orig_model.backbone2.conv2.out_channels * SPARSITY)
assert model.backbone2.fc1.in_features == int(
orig_model.backbone2.fc1.in_features * SPARSITY)