correct

flax_wav2vec2/run_pretraining_loss.py

@@ -9,6 +9,7 @@ from fairseq.criterions.wav2vec_criterion import Wav2VecCriterionConfig, Wav2vec
 from fairseq.tasks.audio_pretraining import AudioPretrainingConfig, AudioPretrainingTask
 
 from transformers import Wav2Vec2ForPreTraining, Wav2Vec2FeatureExtractor
+from transformers.models.wav2vec2.modeling_wav2vec2 import _compute_mask_indices
 
 hf_path = str(sys.argv[1])
 fairseq_wav2vec2_path = str(sys.argv[2])
@@ -16,14 +17,15 @@ fairseq_wav2vec2_path = str(sys.argv[2])
 model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task([fairseq_wav2vec2_path])
 
 feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(hf_path, do_normalize=False)
-hf_model = Wav2Vec2ForPreTraining.from_pretrained(hf_path).train()
+hf_model = Wav2Vec2ForPreTraining.from_pretrained(hf_path, diversity_loss_weight=0.0).train()
 
 model = model[0]
+# set those to 0.0 in the original fairseq code model code
+# also make sure that numpy uses same random seed
 model.cfg["attention_dropout"] = 0.0
 model.cfg["dropout_input"] = 0.0
 model.cfg["dropout_features"] = 0.0
 model.train()
-print(model)
 
 dummy_speech_data = datasets.load_dataset("patrickvonplaten/librispeech_asr_dummy", "clean", split="validation")
 
@@ -42,7 +44,8 @@ attention_mask = inputs.attention_mask
 
 audio_cfg = AudioPretrainingConfig(labels="ltr", data="./data")
 task = AudioPretrainingTask.setup_task(audio_cfg)
-criterion = Wav2vecCriterion(Wav2VecCriterionConfig(infonce=True, log_keys=["prob_perplexity", "code_perplexity", "temp"], loss_weights=[0.1, 10]), task)
+#criterion = Wav2vecCriterion(Wav2VecCriterionConfig(infonce=True, log_keys=["prob_perplexity", "code_perplexity", "temp"]), task, loss_weights=[0.1, 0.0])
+criterion = Wav2vecCriterion(Wav2VecCriterionConfig(infonce=True, log_keys=["prob_perplexity", "code_perplexity", "temp"]), task, loss_weights=[0.0, 0.0])
 
 sample = {
     "net_input": {
@@ -53,9 +56,35 @@ sample = {
 }
 
 torch.manual_seed(0)
-loss, sample_size, log, result = criterion(model, sample)
+loss, sample_size, log = criterion(model, sample)
+
+if attention_mask is not None:
+    mask_indices_seq_length = hf_model._get_feat_extract_output_lengths(input_values.shape[-1])
+    batch_size = input_values.shape[0]
+    # compute real output lengths according to convolution formula
+    output_lengths = hf_model._get_feat_extract_output_lengths(attention_mask.sum(-1)).to(
+        torch.long
+    )
+    sub_attention_mask = torch.zeros(
+        (batch_size, mask_indices_seq_length), dtype=torch.long, device=input_values.device
+    )
+    # these two operations makes sure that all values
+    # before the output lengths indices are attended to
+    sub_attention_mask[
+        (torch.arange(sub_attention_mask.shape[0], device=input_values.device), output_lengths - 1)
+    ] = 1
+    sub_attention_mask = sub_attention_mask.flip([-1]).cumsum(-1).flip([-1]).bool()
+# sample randomly masked indices
+    mask_time_indices = _compute_mask_indices(
+        (batch_size, mask_indices_seq_length),
+        hf_model.config.mask_time_prob,
+        hf_model.config.mask_time_length,
+        attention_mask=sub_attention_mask,
+    )
+    mask_time_indices = torch.tensor(mask_time_indices, device=input_values.device)
+
 torch.manual_seed(0)
-hf_result = hf_model(input_values, attention_mask=attention_mask, mask_time_indices=result["mask_indices"].detach())
+hf_result = hf_model(input_values, attention_mask=attention_mask, mask_time_indices=mask_time_indices)
 
 loss.backward()
 hf_result.loss.backward()
@@ -63,6 +92,19 @@ hf_result.loss.backward()
 print("Loss diff %", 100 * (loss.detach().item() - hf_result.loss.detach().item()) / hf_result.loss.detach())
 print("Loss diff abs", (loss.detach().item() - hf_result.loss.detach().item()))
 
-print("perplexity diff %", 100 * (hf_result.codevector_perplexity.detach().item() - result["prob_perplexity"].detach().item()) / hf_result.codevector_perplexity.detach())
+def grad_norm(model):
+    total_norm = 0.0
+    for p in model.parameters():
+        if p.grad is not None:
+            param_norm = p.grad.detach().data.norm(2)
+            total_norm += param_norm.item() ** 2
+    total_norm = total_norm ** 0.5
+    return total_norm
+
+
+print("Fsq grad norm", grad_norm(model))
+print("HF grad norm", grad_norm(hf_model))
 
 print("Grad max/min diff first layer 'feature_extractor.conv_layers[0].conv.weight'", (hf_model.wav2vec2.feature_extractor.conv_layers[0].conv.weight.grad - model.feature_extractor.conv_layers[0][0].weight.grad).abs().max())
+
+print("Grad max/min diff first layer 'feature_extractor.conv_layers[-1].conv.weight'", (hf_model.wav2vec2.feature_extractor.conv_layers[-1].conv.weight.grad - model.feature_extractor.conv_layers[-1][0].weight.grad).abs().max())