up

flax_wav2vec2/run_pretraining_loss.py

@@ -16,10 +16,14 @@ 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)
+hf_model = Wav2Vec2ForPreTraining.from_pretrained(hf_path).train()
 
 model = model[0]
-model.eval()
+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")
 
@@ -51,9 +55,14 @@ sample = {
 torch.manual_seed(0)
 loss, sample_size, log, result = criterion(model, sample)
 torch.manual_seed(0)
-hf_result = hf_model(input_values, attention_mask=attention_mask, mask_time_indices=result["mask_indices"], fsq_negs=result["negs"])
+hf_result = hf_model(input_values, attention_mask=attention_mask, mask_time_indices=result["mask_indices"].detach())
+
+loss.backward()
+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())
+
+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())

flax_wav2vec2/run_pretraining_loss_flax.py

@@ -4,6 +4,7 @@ import fairseq
 import torch
 import optax
 import jax.numpy as jnp
+import jax
 from flax.training.common_utils import onehot
 
 import soundfile as sf
@@ -20,10 +21,10 @@ model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task([fairse
 
 feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(hf_path, do_normalize=False)
 flax_hf_model = FlaxWav2Vec2ForPreTraining.from_pretrained(hf_path)
-hf_model = Wav2Vec2ForPreTraining.from_pretrained(hf_path)
+hf_model = Wav2Vec2ForPreTraining.from_pretrained(hf_path).train()
 
 model = model[0]
-model.eval()
+model.train()
 
 dummy_speech_data = datasets.load_dataset("patrickvonplaten/librispeech_asr_dummy", "clean", split="validation")
 
@@ -54,7 +55,7 @@ def compute_contrastive_loss(
     neg_is_pos = jnp.concatenate([jnp.full((1,) + loss_logits.shape[1:], False), neg_is_pos], axis=0)
 
     # make sure incorrectly sampled vectors don't contribute to loss
-    loss_logits = jnp.where(neg_is_pos, -1e9, loss_logits)
+    loss_logits = jnp.where(neg_is_pos, -2**30, loss_logits)
 
     predictions = loss_logits.transpose(2, 1, 0).reshape(-1, loss_logits.shape[0])
     targets = ((1 - mask_time_indices) * -100).transpose(1, 0).flatten()
@@ -88,9 +89,11 @@ sample = {
 torch.manual_seed(0)
 loss, sample_size, log, result = criterion(model, sample)
 torch.manual_seed(0)
-hf_result = hf_model(input_values, attention_mask=attention_mask, mask_time_indices=result["mask_indices"])
+hf_result = hf_model(input_values, attention_mask=attention_mask, mask_time_indices=result["mask_indices"], code_vec_indices=result["code_idxs"])
 
-outputs = flax_hf_model(input_values.numpy(), attention_mask=attention_mask.numpy(), mask_time_indices=result["mask_indices"].numpy())
+print(100 * "=")
+
+outputs = flax_hf_model(input_values.numpy(), attention_mask=attention_mask.numpy(), mask_time_indices=result["mask_indices"].numpy(), train=True, gumbel_rng=jax.random.PRNGKey(0), code_vec_indices=result["code_idxs"])
 
 negative_indices = hf_result["sampled_negative_indices"].detach().numpy()
 num_negatives = 100
@@ -114,3 +117,5 @@ loss = contrastive_loss + diversity_loss_weight * diversity_loss
 
 
 print("Loss diff %", 100 * (hf_result.loss.detach().item() - loss.item()) / loss)
+
+print("perplexity diff %", 100 * (outputs.codevector_perplexity - result["prob_perplexity"].detach().item()) / outputs.codevector_perplexity)