def channel_pruning_auto_mode():
# Load trained MNIST model
model = torch.load(os.path.join('../', 'data', 'mnist_trained_on_GPU.pth'))
# Specify the necessary parameters
greedy_params = GreedySelectionParameters(target_comp_ratio=Decimal(0.8),
num_comp_ratio_candidates=10)
auto_params = ChannelPruningParameters.AutoModeParams(
greedy_params, modules_to_ignore=[model.conv1])
data_loader = mnist_torch_model.DataLoaderMnist(cuda=True,
seed=1,
shuffle=True)
params = ChannelPruningParameters(data_loader=data_loader.train_loader,
num_reconstruction_samples=500,
allow_custom_downsample_ops=True,
mode=ChannelPruningParameters.Mode.auto,
params=auto_params)
# Single call to compress the model
results = ModelCompressor.compress_model(
model,
eval_callback=evaluate_model,
eval_iterations=1000,
input_shape=(1, 1, 28, 28),
compress_scheme=CompressionScheme.channel_pruning,
cost_metric=CostMetric.mac,
parameters=params)
compressed_model, stats = results
print(compressed_model)
print(stats) # Stats object can be pretty-printed easily
def channel_pruning_manual_mode():
# Load a trained MNIST model
model = torch.load(os.path.join('../', 'data', 'mnist_trained_on_GPU.pth'))
# Specify the necessary parameters
manual_params = ChannelPruningParameters.ManualModeParams(
[ModuleCompRatioPair(model.conv2, 0.4)])
data_loader = mnist_torch_model.DataLoaderMnist(cuda=True,
seed=1,
shuffle=True)
params = ChannelPruningParameters(
data_loader=data_loader.train_loader,
num_reconstruction_samples=500,
allow_custom_downsample_ops=True,
mode=ChannelPruningParameters.Mode.manual,
params=manual_params)
# Single call to compress the model
results = ModelCompressor.compress_model(
model,
eval_callback=evaluate_model,
eval_iterations=1000,
input_shape=(1, 1, 28, 28),
compress_scheme=CompressionScheme.channel_pruning,
cost_metric=CostMetric.mac,
parameters=params)
compressed_model, stats = results
print(compressed_model)
print(stats) # Stats object can be pretty-printed easily
def weight_svd_manual_mode():
# Load a trained MNIST model
model = torch.load(os.path.join('../', 'data', 'mnist_trained_on_GPU.pth'))
# Specify the necessary parameters
manual_params = WeightSvdParameters.ManualModeParams([
ModuleCompRatioPair(model.conv1, 0.5),
ModuleCompRatioPair(model.conv2, 0.4)
])
params = WeightSvdParameters(mode=WeightSvdParameters.Mode.manual,
params=manual_params,
multiplicity=8)
# Single call to compress the model
results = ModelCompressor.compress_model(
model,
eval_callback=evaluate_model,
eval_iterations=1000,
input_shape=(1, 1, 28, 28),
compress_scheme=CompressionScheme.weight_svd,
cost_metric=CostMetric.mac,
parameters=params)
compressed_model, stats = results
print(compressed_model)
print(stats) # Stats object can be pretty-printed easily
def weight_svd_auto_mode():
# Load trained MNIST model
model = torch.load(os.path.join('../', 'data', 'mnist_trained_on_GPU.pth'))
# Specify the necessary parameters
greedy_params = GreedySelectionParameters(target_comp_ratio=Decimal(0.8),
num_comp_ratio_candidates=10)
rank_select = RankSelectScheme.greedy
auto_params = WeightSvdParameters.AutoModeParams(
rank_select_scheme=rank_select,
select_params=greedy_params,
modules_to_ignore=[model.conv1])
params = WeightSvdParameters(mode=WeightSvdParameters.Mode.auto,
params=auto_params)
# Single call to compress the model
results = ModelCompressor.compress_model(
model,
eval_callback=evaluate_model,
eval_iterations=1000,
input_shape=(1, 1, 28, 28),
compress_scheme=CompressionScheme.weight_svd,
cost_metric=CostMetric.mac,
parameters=params)
compressed_model, stats = results
print(compressed_model)
print(stats) # Stats object can be pretty-printed easily
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_spatial_svd_with_fine_tuning(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)
# load trained MNIST model
data_loader = DataLoaderMnist(cuda=True, seed=1, shuffle=False, train_batch_size=64, test_batch_size=100)
input_shape = (1, 1, 28, 28)
model = torch.load(os.path.join('./', 'data', 'mnist_trained_on_GPU.pth'))
modules_to_ignore = [model.conv1]
greedy_params = aimet_common.defs.GreedySelectionParameters(target_comp_ratio=Decimal(0.8),
num_comp_ratio_candidates=10,
use_monotonic_fit=True)
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=1)
results = ModelCompressor.compress_model(model, mnist_model.evaluate, 10, input_shape,
aimet_common.defs.CompressionScheme.spatial_svd,
cost_metric=aimet_common.defs.CostMetric.mac, parameters=params,
trainer=Trainer(), visualization_url=None)
compressed_model, stats = results
print(compressed_model)
print(stats)
self.assertTrue(math.isclose(float(stats.mac_compression_ratio), 0.7, abs_tol=0.1))
def test_weight_svd_compress_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)
input_shape = (1, 3, 224, 224)
model = models.resnet18(pretrained=True).to(torch.device('cuda'))
manual_params = aimet_torch.defs.WeightSvdParameters.ManualModeParams(
[ModuleCompRatioPair(model.layer1[0].conv1, 0.5),
ModuleCompRatioPair(model.layer2[1].conv2, 0.4), ModuleCompRatioPair(model.fc, 0.4)])
params = aimet_torch.defs.WeightSvdParameters(aimet_torch.defs.WeightSvdParameters.Mode.manual, manual_params)
results = ModelCompressor.compress_model(model, evaluate, 10, input_shape,
aimet_common.defs.CompressionScheme.weight_svd,
cost_metric=aimet_common.defs.CostMetric.mac, parameters=params,
visualization_url=None)
compressed_model, stats = results
print(compressed_model)
print(stats)
self.assertTrue(isinstance(compressed_model.layer1[0].conv1, torch.nn.Sequential))
self.assertEqual(math.floor(64 * 64 * 3 * 3 * 0.5 / (64 + 64 * 3 * 3)),
compressed_model.layer1[0].conv1[0].out_channels)
self.assertTrue(isinstance(compressed_model.fc, torch.nn.Sequential))
self.assertEqual(math.floor(512 * 1000 * 0.4 / (512 + 1000)), compressed_model.fc[0].out_features)
def test_spatial_svd_compress_auto_multi_input_model(self):
torch.manual_seed(1)
numpy.random.seed(1)
AimetLogger.set_level_for_all_areas(logging.DEBUG)
input_shape = [(1, 1, 28, 28), (1, 1, 28, 28)]
model = ModelWithTwoInputs()
greedy_params = aimet_common.defs.GreedySelectionParameters(target_comp_ratio=Decimal(0.65),
num_comp_ratio_candidates=4)
auto_params = aimet_torch.defs.SpatialSvdParameters.AutoModeParams(greedy_params)
params = aimet_torch.defs.SpatialSvdParameters(aimet_torch.defs.SpatialSvdParameters.Mode.auto, auto_params)
mock_eval = FakeEvaluator(input_shape)
mock_eval.return_values = [0.75, 0.50, 0.25, 0.75, 0.50, 0.25, 0.75, 0.50, 0.25,
0.50, 0.50]
results = ModelCompressor.compress_model(model=model, eval_callback=mock_eval, 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)
compressed_model, stats = results
print(compressed_model)
print(stats)
self.assertEqual(8, compressed_model.conv2[0].out_channels)
self.assertEqual((5, 1), compressed_model.conv2[0].kernel_size)
self.assertTrue(math.isclose(float(stats.mac_compression_ratio), 0.25, abs_tol=0.01))
def test_pickled_compression_ratios(self):
visualization_url, process = start_bokeh_server_session(8002)
torch.cuda.empty_cache()
torch.manual_seed(1)
np.random.seed(1)
torch.backends.cudnn.deterministic = True
input_shape = (1, 3, 224, 224)
model = models.resnet18(pretrained=True).to(torch.device('cuda'))
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)
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)
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/greedy_selection_eval_scores_dict.pkl'
self.assertTrue(os.path.exists(comp_ratios_file_path))
layer_comp_ratio_list = CompressionAlgo.unpickle_comp_ratios_list(comp_ratios_file_path)
self.assertEqual(type(layer_comp_ratio_list), list)
compression_visualizations = VisualizeCompression(visualization_url)
compression_visualizations.display_eval_scores(eval_scores_path)
compression_visualizations.display_comp_ratio_plot(comp_ratios_file_path)
os.killpg(os.getpgid(process.pid), signal.SIGTERM)
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_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 test_spatial_svd_compress_auto_with_high_multiplicity(self):
torch.manual_seed(1)
numpy.random.seed(1)
AimetLogger.set_level_for_all_areas(logging.DEBUG)
input_shape = (1, 3, 224, 224)
model = models.resnet18()
modules_to_ignore = [model.conv1,
model.layer1[0].conv1, model.layer1[0].conv2,
model.layer1[1].conv1, model.layer1[1].conv2,
model.layer2[0].downsample[0],
model.layer3[0].conv1, model.layer3[0].conv2,
model.layer3[1].conv1, model.layer3[1].conv2,
model.layer3[0].downsample[0],
model.layer4[0].conv1, model.layer4[0].conv2,
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=4)
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=64)
mock_eval = FakeEvaluator(input_shape)
mock_eval.return_values = [0.75, 0.50, 0.25, 0.75, 0.50, 0.25, 0.75, 0.50, 0.25, 0.75, 0.50, 0.25,
0.50, 0.50]
results = ModelCompressor.compress_model(model=model, eval_callback=mock_eval, 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)
compressed_model, stats = results
print(compressed_model)
print(stats)
self.assertEqual(64, compressed_model.layer2[0].conv2[0].out_channels)
self.assertEqual((3, 1), compressed_model.layer2[0].conv2[0].kernel_size)
self.assertTrue(math.isclose(float(stats.mac_compression_ratio), 0.87, abs_tol=0.01))
def weight_svd_auto_mode(model, comp_ratio=0.8, retrain=False):
input_shape = coord_dataset.mgrid.shape
# Specify the necessary parameters
greedy_params = GreedySelectionParameters(target_comp_ratio=Decimal(comp_ratio),
num_comp_ratio_candidates=20)
#tar_params = TarRankSelectionParameters(num_rank_indices=2)
#rank_select = RankSelectScheme.tar
rank_select = RankSelectScheme.greedy
auto_params = WeightSvdParameters.AutoModeParams(rank_select_scheme=rank_select,
select_params=greedy_params,
)#modules_to_ignore=[model.conv1])
params = WeightSvdParameters(mode=WeightSvdParameters.Mode.auto,
params=auto_params)
#Single call to compress the model
results = ModelCompressor.compress_model(model,
eval_callback=evaluate_model,
eval_iterations=1,
input_shape=input_shape,
compress_scheme=CompressionScheme.weight_svd,
cost_metric=CostMetric.memory,
parameters=params)
compressed_model, stats = results
# torch.save(compressed_model,
# os.path.join(os.path.join(exp_folder, image_name + '/checkpoints/model_aimet_' + str(comp_ratio) +'.pth')))
#print(compressed_model)
print(stats) # Stats object can be pretty-printed easily
# print(os.path.join(os.path.join(exp_folder, image_name + '/checkpoints/model_aimet_.pth')))
#res = check_metrics(dataloader, compressed_model, image_resolution)
#print(res)
loss_fn = partial(loss_functions.image_mse, None)
if retrain:
compressed_model = retrain_model(compressed_model, dataloader,2000, loss_fn, 0.00005, TRAINING_FLAGS['l1_reg'])
torch.save(compressed_model,
os.path.join(
os.path.join(exp_folder, image_name + '/checkpoints/model_aimet_' + str(comp_ratio) + '_retrained.pth')))
res = check_metrics(dataloader, compressed_model, image_resolution)
print(res)
return compressed_model
def test_weight_svd_compress_auto_high_multiplicity(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)
input_shape = (1, 3, 224, 224)
model = models.resnet18(pretrained=True).to(torch.device('cuda'))
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,
model.fc
]
greedy_params = aimet_common.defs.GreedySelectionParameters(target_comp_ratio=Decimal(0.7),
num_comp_ratio_candidates=4,
saved_eval_scores_dict=
'./data/resnet18_eval_scores.pkl')
rank_select = RankSelectScheme.greedy
auto_params = aimet_torch.defs.WeightSvdParameters.AutoModeParams(rank_select_scheme=rank_select,
select_params=greedy_params,
modules_to_ignore=modules_to_ignore)
params = aimet_torch.defs.WeightSvdParameters(aimet_torch.defs.WeightSvdParameters.Mode.auto, auto_params,
multiplicity=64)
results = ModelCompressor.compress_model(model, evaluate, 10, input_shape,
aimet_common.defs.CompressionScheme.weight_svd,
cost_metric=aimet_common.defs.CostMetric.mac, parameters=params,
visualization_url=None)
compressed_model, stats = results
print(compressed_model)
print(stats)
self.assertFalse(isinstance(compressed_model.conv1, torch.nn.Sequential))
self.assertFalse(isinstance(compressed_model.fc, torch.nn.Sequential))
def test_spatial_svd_compress_manual(self):
torch.cuda.empty_cache()
torch.manual_seed(1)
numpy.random.seed(1)
AimetLogger.set_level_for_all_areas(logging.DEBUG)
input_shape = (1, 3, 224, 224)
model = models.resnet18()
manual_params = aimet_torch.defs.SpatialSvdParameters.ManualModeParams(
[ModuleCompRatioPair(model.layer1[0].conv1, 0.5),
ModuleCompRatioPair(model.layer2[1].conv2, 0.4)])
params = aimet_torch.defs.SpatialSvdParameters(aimet_torch.defs.SpatialSvdParameters.Mode.manual, manual_params,
multiplicity=8)
# Only used in this test for baseline and final accuracy - essentially a don't care
mock_eval = MagicMock()
mock_eval.return_value = 50
results = ModelCompressor.compress_model(model=model, eval_callback=mock_eval, 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)
compressed_model, stats = results
print(compressed_model)
print(stats)
# Check that indeed weight svd was applied to some layer
self.assertTrue(isinstance(compressed_model.layer1[0].conv1, torch.nn.Sequential))
self.assertEqual(48, compressed_model.layer1[0].conv1[0].out_channels)
self.assertEqual((3, 1), compressed_model.layer1[0].conv1[0].kernel_size)
self.assertEqual((1, 0), compressed_model.layer1[0].conv1[0].padding)
self.assertTrue(isinstance(compressed_model.layer2[1].conv2, torch.nn.Sequential))
self.assertEqual(80, compressed_model.layer2[1].conv2[0].out_channels)