Examples/torch/compression/channel_pruning.py
def aimet_channel_pruning(model: torch.nn.Module, evaluator: aimet_common.defs.EvalFunction,
                          data_loader: torch_data.DataLoader) -> Tuple[torch.nn.Module, aimet_common.defs.CompressionStats]:
    """
    Compresses the model using AIMET's channel pruning feature

    :param model: The model to compress
    :param evaluator: Evaluator used during compression
    :param dataloader: DataLoader used during compression
    :return: A tuple of compressed model and its statistics
    """

    # Configure the greedy comp-ratio selection algorithm
    greedy_params = aimet_torch.defs.GreedySelectionParameters(target_comp_ratio=Decimal(0.5),
                                                               num_comp_ratio_candidates=10)

    # Configure the auto mode compression. Ignore the first layer of the model (model.conv1).
    auto_params = aimet_torch.defs.ChannelPruningParameters.AutoModeParams(greedy_params,
                                                                           modules_to_ignore=[model.conv1])

    # Configure the parameters for channel pruning compression
    # 50000 reconstruction samples will give better results and is recommended; however we use 5000 here as an example.
    params = aimet_torch.defs.ChannelPruningParameters(data_loader=data_loader,
                                                       num_reconstruction_samples=5000,
                                                       allow_custom_downsample_ops=False,
                                                       mode=aimet_torch.defs.ChannelPruningParameters.Mode.auto,
                                                       params=auto_params)

    scheme = aimet_common.defs.CompressionScheme.channel_pruning  # spatial_svd, weight_svd or channel_pruning
    metric = aimet_common.defs.CostMetric.mac  # mac or memory

    results = ModelCompressor.compress_model(model=model,
                                             eval_callback=evaluator,
                                             eval_iterations=10,
                                             input_shape=(1, 3, 224, 224),
                                             compress_scheme=scheme,
                                             cost_metric=metric,
                                             parameters=params)
    return results

    Compresses the model using AIMET's channel pruning feature

    :param model: The model to compress
    :param evaluator: Evaluator used during compression
    :param dataloader: DataLoader used during compression
    :return: A tuple of compressed model and its statistics
    
Compresses the model using AIMET's channel pruning feature
Compresses the model using AIMET ' s channel pruning feature
def aimet_channel_pruning(model: torch.nn.Module, evaluator: aimet_common.defs.EvalFunction,
                          data_loader: torch_data.DataLoader) -> Tuple[torch.nn.Module, aimet_common.defs.CompressionStats]:
    greedy_params = aimet_torch.defs.GreedySelectionParameters(target_comp_ratio=Decimal(0.5),
                                                               num_comp_ratio_candidates=10)
    auto_params = aimet_torch.defs.ChannelPruningParameters.AutoModeParams(greedy_params,
                                                                           modules_to_ignore=[model.conv1])
    params = aimet_torch.defs.ChannelPruningParameters(data_loader=data_loader,
                                                       num_reconstruction_samples=5000,
                                                       allow_custom_downsample_ops=False,
                                                       mode=aimet_torch.defs.ChannelPruningParameters.Mode.auto,
                                                       params=auto_params)
    scheme = aimet_common.defs.CompressionScheme.channel_pruning  
    metric = aimet_common.defs.CostMetric.mac  
    results = ModelCompressor.compress_model(model=model,
                                             eval_callback=evaluator,
                                             eval_iterations=10,
                                             input_shape=(1, 3, 224, 224),
                                             compress_scheme=scheme,
                                             cost_metric=metric,
                                             parameters=params)
    return results
def aimet_channel_pruning ( model : torch . nn . Module , evaluator : aimet_common . defs . EvalFunction , data_loader : torch_data . DataLoader ) -> Tuple [ torch . nn . Module , aimet_common . defs . CompressionStats ] : greedy_params = aimet_torch . defs . GreedySelectionParameters ( target_comp_ratio = Decimal ( 0.5 ) , num_comp_ratio_candidates = 10 ) auto_params = aimet_torch . defs . ChannelPruningParameters . AutoModeParams ( greedy_params , modules_to_ignore = [ model . conv1 ] ) params = aimet_torch . defs . ChannelPruningParameters ( data_loader = data_loader , num_reconstruction_samples = 5000 , allow_custom_downsample_ops = False , mode = aimet_torch . defs . ChannelPruningParameters . Mode . auto , params = auto_params ) scheme = aimet_common . defs . CompressionScheme . channel_pruning metric = aimet_common . defs . CostMetric . mac results = ModelCompressor . compress_model ( model = model , eval_callback = evaluator , eval_iterations = 10 , input_shape = ( 1 , 3 , 224 , 224 ) , compress_scheme = scheme , cost_metric = metric , parameters = params ) return results
Compresses the model using AIMET's channel pruning feature
Compresses the model using AIMET ' s channel pruning feature
"""
    Compresses the model using AIMET's channel pruning feature

    :param model: The model to compress
    :param evaluator: Evaluator used during compression
    :param dataloader: DataLoader used during compression
    :return: A tuple of compressed model and its statistics
    """ # Configure the greedy comp-ratio selection algorithm # Configure the auto mode compression. Ignore the first layer of the model (model.conv1). # Configure the parameters for channel pruning compression # 50000 reconstruction samples will give better results and is recommended; however we use 5000 here as an example. # spatial_svd, weight_svd or channel_pruning # mac or memory
A tuple of compressed model and its statistics
A tuple of compressed model and its statistics
model
torch.nn.Module
The model to compress
The model to compress
evaluator
aimet_common.defs.EvalFunction
Evaluator used during compression
Evaluator used during compression
data_loader
torch_data.DataLoader
dataloader
DataLoader used during compression
DataLoader used during compression