def equalize_model(model: torch.nn.Module, input_shapes: Union[Tuple, List[Tuple]]): """ High-level API to perform Cross-Layer Equalization (CLE) on the given model. The model is equalized in place. :param model: Model to equalize :param input_shapes: Shape of the input (can be a tuple or a list of tuples if multiple inputs) :return: None """ device = get_device(model) model.cpu() # fold batchnorm layers folded_pairs = fold_all_batch_norms(model, input_shapes) bn_dict = {} for conv_bn in folded_pairs: bn_dict[conv_bn[0]] = conv_bn[1] # replace any ReLU6 layers with ReLU utils.replace_modules_of_type1_with_type2(model, torch.nn.ReLU6, torch.nn.ReLU) # perform cross-layer scaling on applicable layer sets cls_set_info_list = CrossLayerScaling.scale_model(model, input_shapes) # high-bias fold HighBiasFold.bias_fold(cls_set_info_list, bn_dict) model.to(device=device)
def main(): args = arguments() seed(args) model = DeepLab(backbone='mobilenet', output_stride=16, num_classes=21, sync_bn=False) model.eval() from aimet_torch import batch_norm_fold from aimet_torch import utils args.input_shape = (1, 3, 513, 513) batch_norm_fold.fold_all_batch_norms(model, args.input_shape) utils.replace_modules_of_type1_with_type2(model, torch.nn.ReLU6, torch.nn.ReLU) if args.checkpoint_path: model.load_state_dict(torch.load(args.checkpoint_path)) else: raise ValueError('checkpoint path {} must be specified'.format( args.checkpoint_path)) data_loader_kwargs = {'worker_init_fn': work_init, 'num_workers': 0} train_loader, val_loader, test_loader, num_class = make_data_loader( args, **data_loader_kwargs) eval_func_quant = model_eval(args, val_loader) eval_func = model_eval(args, val_loader) from aimet_common.defs import QuantScheme from aimet_torch.quantsim import QuantizationSimModel if hasattr(args, 'quant_scheme'): if args.quant_scheme == 'range_learning_tf': quant_scheme = QuantScheme.training_range_learning_with_tf_init elif args.quant_scheme == 'range_learning_tfe': quant_scheme = QuantScheme.training_range_learning_with_tf_enhanced_init elif args.quant_scheme == 'tf': quant_scheme = QuantScheme.post_training_tf elif args.quant_scheme == 'tf_enhanced': quant_scheme = QuantScheme.post_training_tf_enhanced else: raise ValueError("Got unrecognized quant_scheme: " + args.quant_scheme) kwargs = { 'quant_scheme': quant_scheme, 'default_param_bw': args.default_param_bw, 'default_output_bw': args.default_output_bw, 'config_file': args.config_file } print(kwargs) sim = QuantizationSimModel(model.cpu(), input_shapes=args.input_shape, **kwargs) sim.compute_encodings(eval_func_quant, (1024, True)) post_quant_top1 = eval_func(sim.model.cuda(), (99999999, True)) print("Post Quant mIoU :", post_quant_top1)
def test_replace_relu_with_relu6(self): model = torchvision.models.resnet18() model.eval() utils.replace_modules_of_type1_with_type2(model, torch.nn.ReLU, torch.nn.ReLU6) # check - no ReLU modules left in the model anymore for module in model.modules(): self.assertTrue(not isinstance(module, torch.nn.ReLU)) # sanity-check: forward pass continues to work with torch.no_grad(): x = torch.rand(1, 3, 224, 224) output = model(x)
def cross_layer_equalization_manual(): model = models.resnet18(pretrained=True) model = model.eval() # Batch Norm Fold # Create a list of conv/linear and BN layers for folding forward or backward layer_list = [(model.conv1, model.bn1), (model.layer1[0].conv1, model.layer1[0].bn1)] # Save the corresponding BN layers (needed only for high bias folding) bn_dict = {} for conv_bn in layer_list: bn_dict[conv_bn[0]] = conv_bn[1] batch_norm_fold.fold_given_batch_norms(model, layer_list) # Replace any ReLU6 layers with ReLU utils.replace_modules_of_type1_with_type2(model, torch.nn.ReLU6, torch.nn.ReLU) # Cross Layer Scaling # Create a list of consecutive conv layers to be equalized consecutive_layer_list = [(model.conv1, model.layer1[0].conv1), (model.layer1[0].conv1, model.layer1[0].conv2)] scaling_factor_list = cross_layer_equalization.CrossLayerScaling.scale_cls_sets( consecutive_layer_list) # High Bias Fold # Create a list of consecutive conv layers whose previous layers bias has to be folded to next layers bias ClsSetInfo = cross_layer_equalization.ClsSetInfo ClsPairInfo = cross_layer_equalization.ClsSetInfo.ClsSetLayerPairInfo cls_set_info_list = [ ClsSetInfo( ClsPairInfo(model.conv1, model.layer1[0].conv1, scaling_factor_list[0], True)), ClsSetInfo( ClsPairInfo(model.layer1[0].conv1, model.layer1[0].conv2, scaling_factor_list[1], True)) ] cross_layer_equalization.HighBiasFold.bias_fold(cls_set_info_list, bn_dict)
def cross_layer_equalization_depthwise_layers(): model = MobileNetV2().to(torch.device('cpu')) model.eval() # Batch Norm Fold # Create a list of conv/linear and BN layers for folding forward or backward layer_list = [(model.features[0][0], model.features[0][1]), (model.features[1].conv[0], model.features[1].conv[1]), (model.features[1].conv[3], model.features[1].conv[4])] # Save the corresponding BN layers (needed only for high bias folding) bn_dict = {} for conv_bn in layer_list: bn_dict[conv_bn[0]] = conv_bn[1] batch_norm_fold.fold_given_batch_norms(model, layer_list) # Replace any ReLU6 layers with ReLU utils.replace_modules_of_type1_with_type2(model, torch.nn.ReLU6, torch.nn.ReLU) # Cross Layer Scaling # Create a list of consecutive conv layers to be equalized consecutive_layer_list = [(model.features[0][0], model.features[1].conv[0], model.features[1].conv[3])] scaling_factor_list = cross_layer_equalization.CrossLayerScaling.scale_cls_sets( consecutive_layer_list) # High Bias Fold # Create a list of consecutive conv layers whose previous layers bias has to be folded to next layers bias ClsSetInfo = cross_layer_equalization.ClsSetInfo ClsPairInfo = cross_layer_equalization.ClsSetInfo.ClsSetLayerPairInfo cls_set_info_list = [ ClsSetInfo( ClsPairInfo(model.features[0][0], model.features[1].conv[0], scaling_factor_list[0][0], True)), ClsSetInfo( ClsPairInfo(model.features[1].conv[0], model.features[1].conv[3], scaling_factor_list[0][1], True)) ] cross_layer_equalization.HighBiasFold.bias_fold(cls_set_info_list, bn_dict)
def cross_layer_equalization_auto_step_by_step(): model = models.resnet18(pretrained=True) model = model.eval() input_shape = (1, 3, 224, 224) # Fold batchnorm layers folded_pairs = batch_norm_fold.fold_all_batch_norms(model, input_shape) bn_dict = {} for conv_bn in folded_pairs: bn_dict[conv_bn[0]] = conv_bn[1] # Replace any ReLU6 layers with ReLU utils.replace_modules_of_type1_with_type2(model, torch.nn.ReLU6, torch.nn.ReLU) # Perform cross-layer scaling on applicable layer sets cls_set_info_list = cross_layer_equalization.CrossLayerScaling.scale_model( model, input_shape) # Perform high-bias fold cross_layer_equalization.HighBiasFold.bias_fold(cls_set_info_list, bn_dict)