def test_channel_pruning_manual(self):
torch.cuda.empty_cache()
torch.manual_seed(1)
numpy.random.seed(1)
torch.backends.cudnn.deterministic = True
AimetLogger.set_level_for_all_areas(logging.DEBUG)
data_loader = ImageNetDataLoader(image_dir, image_size, batch_size, num_workers)
input_shape = (1, 3, 224, 224)
model = models.resnet18(pretrained=True).to(torch.device('cuda'))
manual_params = ChannelPruningParameters.ManualModeParams([ModuleCompRatioPair(model.layer1[0].conv2, 0.3),
ModuleCompRatioPair(model.layer2[1].conv1, 0.5)])
params = ChannelPruningParameters(data_loader.train_loader, 5000,
True,
aimet_torch.defs.ChannelPruningParameters.Mode.manual,
manual_params, multiplicity=8)
compressed_model, stats = ModelCompressor.compress_model(model, evaluate, 10, input_shape,
aimet_common.defs.CompressionScheme.channel_pruning,
cost_metric=aimet_common.defs.CostMetric.mac,
parameters=params, visualization_url=None)
baseline_model_accuracy = stats.baseline_model_accuracy
compressed_best_model_accuracy = stats.compressed_model_accuracy
self.assertTrue(baseline_model_accuracy >= compressed_best_model_accuracy)
self.assertEqual(24, compressed_model.layer1[0].conv2.in_channels)
def test_channel_pruning_compress_auto_resnet(self):
torch.cuda.empty_cache()
torch.manual_seed(1)
numpy.random.seed(1)
torch.backends.cudnn.deterministic = True
AimetLogger.set_level_for_all_areas(logging.DEBUG)
data_loader = ImageNetDataLoader(image_dir, image_size, batch_size, num_workers)
input_shape = (1, 3, 224, 224)
model = models.resnet18(pretrained=False).to(torch.device('cuda'))
model.eval()
modules_to_ignore = [model.conv1,
model.layer2[0].downsample[0],
model.layer3[0].downsample[0],
model.layer4[0].downsample[0],
model.layer4[1].conv1,
model.layer4[1].conv2
]
greedy_params = aimet_common.defs.GreedySelectionParameters(target_comp_ratio=Decimal(0.65),
num_comp_ratio_candidates=10,
use_monotonic_fit=True,
saved_eval_scores_dict=
'./data/resnet18_eval_scores.pkl')
auto_params = ChannelPruningParameters.AutoModeParams(greedy_params,
modules_to_ignore=modules_to_ignore)
# selecting single batch for reconstruction
# num_reconstruction_samples = 50
# 50 / 10 (samples_per_image) = 5 = batch size
params = ChannelPruningParameters(data_loader=data_loader.train_loader,
num_reconstruction_samples=50,
allow_custom_downsample_ops=True,
mode=aimet_torch.defs.ChannelPruningParameters.Mode.auto,
params=auto_params, multiplicity=8)
results = ModelCompressor.compress_model(model=model, eval_callback=evaluate, eval_iterations=5,
input_shape=input_shape,
compress_scheme=aimet_common.defs.CompressionScheme.channel_pruning,
cost_metric=aimet_common.defs.CostMetric.mac, parameters=params,
visualization_url=None)
compressed_model, stats = results
print(compressed_model)
print(stats)
self.assertNotEqual(model, compressed_model)
self.assertTrue(0.6 < float(stats.mac_compression_ratio) < 0.65)
def model_eval(model, early_stopping_iterations, use_cuda):
"""
:param model: model to be evaluated
:param early_stopping_iterations: if None, data loader will iterate over entire validation data
:return: top_1_accuracy on validation data
"""
data_loader = ImageNetDataLoader(image_dir, image_size, batch_size, num_workers)
if early_stopping_iterations is not None:
# wrapper around validation data loader to run only 'X' iterations to save time
val_loader = IterFirstX(data_loader.val_loader, early_stopping_iterations)
else:
# iterate over entire validation data set
val_loader = data_loader.val_loader
criterion = nn.CrossEntropyLoss().cuda()
evaluator = create_stand_alone_supervised_classification_evaluator(model, criterion, use_cuda)
evaluator.run(val_loader)
return evaluator.state.metrics['top_1_accuracy']
def model_compression_with_visualization():
"""
Code example for compressing a model with a visualization url provided.
"""
visualization_url, process = start_bokeh_server_session(8002)
ImageNetDataLoader(image_dir, image_size, batch_size, num_workers)
input_shape = (1, 3, 224, 224)
model = models.resnet18(pretrained=True).to(torch.device('cuda'))
modules_to_ignore = [model.conv1]
greedy_params = aimet_common.defs.GreedySelectionParameters(
target_comp_ratio=Decimal(0.65),
num_comp_ratio_candidates=10,
saved_eval_scores_dict='../data/resnet18_eval_scores.pkl')
auto_params = aimet_torch.defs.SpatialSvdParameters.AutoModeParams(
greedy_params, modules_to_ignore=modules_to_ignore)
params = aimet_torch.defs.SpatialSvdParameters(
aimet_torch.defs.SpatialSvdParameters.Mode.auto,
auto_params,
multiplicity=8)
# If no visualization URL is provided, during model compression execution no visualizations will be published.
ModelCompressor.compress_model(
model=model,
eval_callback=evaluate,
eval_iterations=5,
input_shape=input_shape,
compress_scheme=aimet_common.defs.CompressionScheme.spatial_svd,
cost_metric=aimet_common.defs.CostMetric.mac,
parameters=params,
visualization_url=None)
comp_ratios_file_path = './data/greedy_selection_comp_ratios_list.pkl'
eval_scores_path = '../data/resnet18_eval_scores.pkl'
# A user can visualize the eval scores dictionary and optimal compression ratios by executing the following code.
compression_visualizations = VisualizeCompression(visualization_url)
compression_visualizations.display_eval_scores(eval_scores_path)
compression_visualizations.display_comp_ratio_plot(comp_ratios_file_path)
def test_bias_correction_hybrid(self):
torch.manual_seed(10)
model = MobileNetV2().to(torch.device('cpu'))
model.eval()
module_prop_list = aimet_torch.bias_correction.find_all_conv_bn_with_activation(
model, input_shape=(1, 3, 224, 224))
batch_norm_fold.fold_all_batch_norms(model, (1, 3, 224, 224))
model_copy = copy.deepcopy(model)
model.eval()
model_copy.eval()
image_dir = './data/tiny-imagenet-200'
image_size = 224
batch_size = 1
num_workers = 1
data_loader = ImageNetDataLoader(image_dir, image_size, batch_size,
num_workers)
params = QuantParams(weight_bw=4,
act_bw=4,
round_mode="nearest",
quant_scheme=QuantScheme.post_training_tf)
bias_correction.correct_bias(model.to(device="cuda"), params, 1,
data_loader.train_loader, 1,
module_prop_list, False)
assert (np.allclose(
model.features[0][0].bias.detach().cpu().numpy(),
model_copy.features[0][0].bias.detach().cpu().numpy()))
assert (np.allclose(
model.features[1].conv[0].bias.detach().cpu().numpy(),
model_copy.features[1].conv[0].bias.detach().cpu().numpy()))
# To check if wrappers got removed
assert (isinstance(model.features[11].conv[0], nn.Conv2d))
def model_train(model, epochs, callback=None):
"""
:param model: model
:param epochs: number of epochs
:return: accuracy after each epoch on training , validation data
"""
data_loader = ImageNetDataLoader(two_class_image_dir, image_size, batch_size, num_workers)
criterion = nn.CrossEntropyLoss().cuda()
lr = 0.01
momentum = 0.9
lr_step_size = 0.01
lr_gamma = 0.01
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=lr_step_size, gamma=lr_gamma)
trainer, evaluator = create_supervised_classification_trainer(model=model, loss_fn=criterion, optimizer=optimizer,
val_loader=data_loader.val_loader,
learning_rate_scheduler=exp_lr_scheduler,
use_cuda=True,
callback=callback)
trainer.run(data_loader.train_loader, max_epochs=epochs)
return trainer.state.metrics['top_1_accuracy'], evaluator.state.metrics['top_1_accuracy']