Grabbing the output for all Convolution layers in Wav2VecForCTC Model
Author: priyammaz
Pull request: no

### Feature request

I want to be able to grab the output of all 7 convolution blocks from the Wav2VecForCTC model but I cant think of a way to do it. I tried to update the forward function of the Wav2Vec2FeatureEncoder with a new attribute that stores each hidden state of the convolution to a list but the moment I load the default pretrained model, the attribute no longer exists. 

### Motivation

I am working on model explainability and the option to grab outputs of each convolution at every step would allow me to do a deeper dive of how the model interprets different phonemes. 

### Your contribution

I am happy to help out how I can! Im not really sure where to even being though with this, maybe there is something simple that I am missing?

Comments:
cc @sanchit-gandhi and @ArthurZucker 

Hey @priyammaz - you can append the hidden-states for each layer to a tuple in the same way that we do for the Wav2Vec2Encoder.

In the forward call of `Wav2Vec2FeatureEncoder`:
```python
    def forward(self, input_values, output_hidden_states=False):
        all_hidden_states = () if output_hidden_states else None
        hidden_states = input_values[:, None]

        # make sure hidden_states require grad for gradient_checkpointing
        if self._requires_grad and self.training:
            hidden_states.requires_grad = True

        for conv_layer in self.conv_layers:
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)
            if self._requires_grad and self.gradient_checkpointing and self.training:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs)

                    return custom_forward

                hidden_states = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(conv_layer),
                    hidden_states,
                )
            else:
                hidden_states = conv_layer(hidden_states)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        return BaseModelOutput(
            last_hidden_state=hidden_states, hidden_states=all_hidden_states,
        )
```

In the forward call of `Wav2Vec2Model`:
```python
    def forward(
        self,
        input_values: Optional[torch.Tensor],
        attention_mask: Optional[torch.Tensor] = None,
        mask_time_indices: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, Wav2Vec2BaseModelOutput]:
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        conv_features = self.feature_extractor(input_values, output_hidden_states=output_hidden_states)
        extract_features = conv_features[0].transpose(1, 2)

        if attention_mask is not None:
            # compute reduced attention_mask corresponding to feature vectors
            attention_mask = self._get_feature_vector_attention_mask(
                extract_features.shape[1], attention_mask, add_adapter=False
            )

        hidden_states, extract_features = self.feature_projection(extract_features)
        hidden_states = self._mask_hidden_states(
            hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask
        )

        encoder_outputs = self.encoder(
            hidden_states,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        hidden_states = encoder_outputs[0]

        if self.adapter is not None:
            hidden_states = self.adapter(hidden_states)

        if not return_dict:
            return (hidden_states, extract_features) + encoder_outputs[1:]

        all_hidden_states = conv_features[1] + encoder_outputs.hidden_states if output_hidden_states else None

        return Wav2Vec2BaseModelOutput(
            last_hidden_state=hidden_states,
            extract_features=extract_features,
            hidden_states=all_hidden_states,
            attentions=encoder_outputs.attentions,
        )
```

All in all, it looks something like this: https://github.com/sanchit-gandhi/codesnippets/blob/main/modeling_wav2vec2_with_conv_states.py


ThanK you so much! I will give this a try this weekend and let you know if I am stuck anywhere, I am still learning the HuggingFace platform!

This worked perfectly thank you so much! 

Cool! Glad to hear that @priyammaz! Thinking about it more, we also apply a feature projection after the last CNN layer:
https://github.com/sanchit-gandhi/codesnippets/blob/cb6a463b2b948a78081b382d51c062ca0ae8de31/modeling_wav2vec2_with_conv_states.py#L1324
This feature projection is essentially just layer norm followed by a linear layer:
https://github.com/sanchit-gandhi/codesnippets/blob/cb6a463b2b948a78081b382d51c062ca0ae8de31/modeling_wav2vec2_with_conv_states.py#L485
You may also want to return the output of this feature projection layer if it's of interest to your research (you can do so simply by appending the outputs to our tuple of `output_hidden_states` as we do for the conv layer outputs:
https://github.com/sanchit-gandhi/codesnippets/blob/cb6a463b2b948a78081b382d51c062ca0ae8de31/modeling_wav2vec2_with_conv_states.py#L1345 

Thanks for the info! I will definitely give that a try!