quic / aimet

AIMET is a library that provides advanced quantization and compression techniques for trained neural network models.
https://quic.github.io/aimet-pages/index.html
Other
2.13k stars 382 forks source link

Do channelprunning support RNN model? #360

Closed hellolingling closed 1 year ago

hellolingling commented 3 years ago

`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)

` As Channel Prunning need to load data, when I test channel prunning on RNN model, the input of the model is [current frame, last frame result by model]. But the code in AIMET\build\staging\universal\lib\python\aimet_torch\data_subsampler.py don't support RNN. How can we change the code to support RNN model for channel prunning? @quic-bharathr @ @quic-akhobare

quic-akhobare commented 3 years ago

Hi @hellolingling - thanks for the question,

We have not tested channel pruning with RNN models so far. The RNN layers itself are not supported for input channel pruning. Could you please tell us a little more information about your intent?

Let me know if the question was not clear.

hellolingling commented 3 years ago

Hi @quic-akhobare - thanks for your reply,

I am intent to remove input channels for conv layers in other part of the model, but the code in AIMET\build\staging\universal\lib\python\aimet_torch\data_subsampler.py showing as follows requires the dataloader should be in [input_data, labels] format, but the input of our RNN model is [input_data, preframe_output_of_model], so how to change the code in data_subsamper.py to apply channel prunning in RNN model?

`class DataSubSampler: """ Utilities to sub-sample data for weight reconstruction """

@staticmethod
def _forward_pass(model: torch.nn.Module, batch: Union[torch.Tensor, List, Tuple]):
    """
    forward pass depending model allocation on CPU / GPU till StopForwardException
    :param model: model
    :param batch: batch
    """
    # keep the model in eval mode
    model.eval()

    # get the model's device placement information
    device = utils.get_device(model)
    # place the batch to appropriate device
    batch = utils.change_tensor_device_placement(batch, device)

    if isinstance(batch, torch.Tensor):
        batch = [batch]

    try:
        with torch.no_grad():
            _ = model(*batch)
    except StopForwardException:
        pass

@classmethod
def get_sub_sampled_data(cls, orig_layer: Union[torch.nn.Conv2d, torch.nn.Linear],
                         pruned_layer: Union[torch.nn.Conv2d, torch.nn.Linear],
                         orig_model: torch.nn.Module, comp_model: torch.nn.Module, data_loader: Iterator,
                         num_reconstruction_samples: int) -> (np.ndarray, np.ndarray):
    # pylint: disable=too-many-locals
    """
    Get all the input data from pruned model and output data from original model

    :param orig_layer: original layer
    :param pruned_layer: pruned layer
    :param orig_model: original model, un-pruned, used to provide the actual outputs
    :param comp_model: comp. model, this is potentially already pruned in the upstreams layers of given layer name
    :param data_loader: data loader
    :param num_reconstruction_samples: The number of reconstruction samples
    :return: input_data, output_data
    """

    if isinstance(orig_layer, torch.nn.Conv2d) and isinstance(pruned_layer, torch.nn.Conv2d):
        sub_sampler = Conv2dSubSampler()

    elif isinstance(orig_layer, torch.nn.Linear) and isinstance(pruned_layer, torch.nn.Linear):
        sub_sampler = LinearSubSampler()

    else:
        raise ValueError('Layer type not supported!')

    # forward pass for given number of batches for both original model and compressed model
    for batch_index, batch in enumerate(data_loader):

        assert isinstance(batch, (tuple, list)), 'data loader should provide data in list or tuple format' \
                                                 '(input_data, labels) or [input_data, labels]'

        batch, _ = batch

        DataSubSampler._forward_pass(orig_model, batch)
        DataSubSampler._forward_pass(comp_model, batch)`
quic-hitameht commented 3 years ago

@hellolingling thanks for sharing your question,

Yes I can see that the 'assert isinstance(batch, (tuple, list))' is very strong and problematic and we will need to change that.

In the meantime. can you make change in your data loader to return tuple - ([input_data, preframe_output_of_model], None)? This way [input_data, preframe_output_of_model] will be passed to your model as it is.

Please let us know if this solution works.

quic-mangal commented 1 year ago

Closing this issue due to inactivity. Please re-open it/ create a new issue if you need further help.