fix bugs for finetuning unispeech

microsoft · Apr 7, 2022 · e3043e2 · e3043e2
1 parent 295c961
commit e3043e2
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/src/fairseq/models/unispeech/unispeech.py b/src/fairseq/models/unispeech/unispeech.py
@@ -63,6 +63,9 @@ def forward(self, **kwargs):
         x = self.w2v_encoder(**kwargs)
         return x
 
+    def remove_pretraining_modules(self):
+        self.w2v_encoder.proj = None
+
 class Wav2VecEncoder(FairseqEncoder):
     def __init__(self, cfg, task):
         super().__init__(task.source_dictionary)

diff --git a/src/fairseq/models/wav2vec/wav2vec2_asr.py b/src/fairseq/models/wav2vec/wav2vec2_asr.py
@@ -353,16 +353,19 @@ def __init__(self, cfg: Wav2Vec2AsrConfig, output_size=None):
         task = tasks.setup_task(w2v_args.task)
         model = task.build_model(w2v_args.model)
 
+        model.remove_pretraining_modules()
         if state is not None and not cfg.no_pretrained_weights:
             model.load_state_dict(state["model"], strict=True)
 
-        model.remove_pretraining_modules()
 
         super().__init__(task.source_dictionary)
 
         d = w2v_args.model.encoder_embed_dim
 
-        self.w2v_model = model
+        if hasattr(model, 'w2v_encoder'):
+            self.w2v_model = model.w2v_encoder.w2v_model
+        else:
+            self.w2v_model = model
 
         self.final_dropout = nn.Dropout(cfg.final_dropout)
         self.freeze_finetune_updates = cfg.freeze_finetune_updates