def test_staged_scheduler_with_empty_quantization(): config = get_squeezenet_quantization_config() config['compression'].update({ 'params': { "activations_quant_start_epoch": 1, "weights_quant_start_epoch": 2, } }) model = squeezenet1_1_custom(num_classes=10, pretrained=False, dropout=0) model, algo = create_compressed_model_and_algo_for_test(model, config) scheduler = algo.scheduler for module in algo.all_quantizations.values(): assert not module.is_enabled_quantization() scheduler.epoch_step() for module in algo.all_quantizations.values(): if module.is_weights: assert not module.is_enabled_quantization() else: assert module.is_enabled_quantization() scheduler.epoch_step() for module in algo.all_quantizations.values(): assert module.is_enabled_quantization()
def test_hawq_on_single_conv_without_quantizers(_seed, dataset_dir, tmp_path, params: HAWQTestParams): config = get_squeezenet_quantization_config(batch_size=params.batch_size) iter_number = params.iter_number tolerance = 4e-4 model = squeezenet1_1_custom(num_classes=10, pretrained=False, dropout=0) from examples.common.models.classification.squeezenet import model_urls load_state(model, model_zoo.load_url(model_urls['squeezenet1_1'])) model = model.cuda() criterion = nn.CrossEntropyLoss().cuda() if not dataset_dir: dataset_dir = str(tmp_path) data_loader, _ = create_test_dataloaders(config.get("model_size"), dataset_dir, params.batch_size) device = next(model.parameters()).device for _, param in model.named_parameters(): param.requires_grad = False first_conv = next(iter(get_all_modules_by_type(model, 'Conv2d').values())) first_conv.weight.requires_grad = True trace_estimator = HessianTraceEstimator(model, criterion, device, data_loader, params.num_data_points) actual_state = trace_estimator.get_average_traces(max_iter=iter_number, tolerance=tolerance) assert math.isclose(actual_state.item(), params.ref_trace, rel_tol=1e-09)
def test_staged_scheduler_with_hawq(): config = get_squeezenet_quantization_config() config['compression'].update({ 'params': { "activations_quant_start_epoch": 1, "weights_quant_start_epoch": 2, }, 'initializer': { 'range': { 'num_init_samples': 1 }, 'precision': { "type": "hawq", "num_data_points": 1, "iter_number": 1, "tolerance": 1 } } }) num_classes = 10 model = squeezenet1_1(num_classes=num_classes, dropout=0) input_infos_list = create_input_infos(config) input_sample_size = input_infos_list[0].shape data_loader = DataLoader( HawqDatasetMock(input_sample_size[1:], num_classes), batch_size=1, num_workers=0, # Workaround for PyTorch MultiprocessingDataLoader issues shuffle=False) criterion = nn.CrossEntropyLoss().cuda() config = register_default_init_args(config, data_loader, criterion) model, algo = create_compressed_model_and_algo_for_test(model, config) scheduler = algo.scheduler for module in algo.all_quantizations.values(): assert not module.is_enabled_quantization() scheduler.epoch_step() for module in algo.all_quantizations.values(): assert not module.is_enabled_quantization() scheduler.epoch_step() for module in algo.all_quantizations.values(): if module.is_weights: assert not module.is_enabled_quantization() else: assert module.is_enabled_quantization() scheduler.epoch_step() for module in algo.all_quantizations.values(): assert module.is_enabled_quantization()
def test_staged_scheduler_with_range_init(): config = get_squeezenet_quantization_config() config['compression'].update({ 'params': { "activations_quant_start_epoch": 1, "weights_quant_start_epoch": 2, }, 'initializer': { 'range': { 'num_init_samples': 1 } } }) model = squeezenet1_1(num_classes=10, dropout=0) input_infos_list = create_input_infos(config) input_sample_size = input_infos_list[0].shape data_loader = DataLoader( OnesDatasetMock(input_sample_size[1:]), batch_size=1, num_workers=0, # Workaround for PyTorch MultiprocessingDataLoader issues shuffle=False) config.register_extra_structs([QuantizationRangeInitArgs(data_loader)]) model, algo = create_compressed_model_and_algo_for_test(model, config) scheduler = algo.scheduler for module in algo.all_quantizations.values(): assert not module.is_enabled_quantization() scheduler.epoch_step() for module in algo.all_quantizations.values(): assert not module.is_enabled_quantization() scheduler.epoch_step() for module in algo.all_quantizations.values(): if module.is_weights: assert not module.is_enabled_quantization() else: assert module.is_enabled_quantization() scheduler.epoch_step() for module in algo.all_quantizations.values(): assert module.is_enabled_quantization()
def create_hawq_test_config(batch_size, num_data_points): config = get_squeezenet_quantization_config() config['batch_size'] = batch_size config['compression'].update({ 'initializer': { 'precision': { "type": "hawq", "bits": [ 4, 8, 6, 7, 5 ], "num_data_points": num_data_points, "iter_number": 1, "tolerance": 1e-2 }, 'range': { 'num_init_steps': 1 } }}) return config