Skip to content

Output hiddens states from all the transformer layers of huggingface_wav2vec #1570

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Adel-Moumen merged 2 commits into from
Sep 17, 2022
Merged

Output hiddens states from all the transformer layers of huggingface_wav2vec #1570

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a3d4ffe
output all hiddens states
BenoitWang Sep 9, 2022
bc35235
minor fix
BenoitWang Sep 17, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 17 additions & 2 deletions speechbrain/lobes/models/huggingface_wav2vec.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,12 @@ class HuggingFaceWav2Vec2(nn.Module):
If True, the model will apply spec augment on the output of feature extractor
(inside huggingface Wav2VecModel() class).
If False, the model will not apply spec augment. We set this to false to prevent from doing it twice.
output_all_hiddens : bool (default: False)
If True, the forward function outputs the hidden states from all transformer layers.
For example wav2vec2-base has 12 transformer layers and the output is of shape (13, B, T, C),
where a projection of the CNN output is added to the beginning.
If False, the forward function outputs the hidden states only from the last transformer layer.

Example
-------
>>> inputs = torch.rand([10, 600])
Expand All @@ -100,6 +106,7 @@ def __init__(
freeze=True,
freeze_feature_extractor=False,
apply_spec_augment=False,
output_all_hiddens=False,
):
super().__init__()

Expand Down Expand Up @@ -145,6 +152,7 @@ def __init__(
self.model.train()
if self.freeze_feature_extractor:
self.model.feature_extractor._freeze_parameters()
self.output_all_hiddens = output_all_hiddens

def _from_pretrained(self, source, config, model, save_path):
"""This function manages the source checking and loading of the params.
Expand Down Expand Up @@ -278,11 +286,18 @@ def extract_features(self, wav):
wav = F.layer_norm(wav, wav.shape)

# Extract wav2vec output
out = self.model(wav)[0]
out = self.model(wav, output_hidden_states=True)

if self.output_all_hiddens:
out = torch.stack(list(out[2]), dim=0)
norm_shape = out.shape[-3:]
else:
out = out[0]
norm_shape = out.shape

# We normalize the output if required
if self.output_norm:
out = F.layer_norm(out, out.shape)
out = F.layer_norm(out, norm_shape)

return out

Expand Down