feat: implement mixed precision training and gradient accumulation in training loop

app/training/wav2vec2_classifier.py

@@ -241,26 +241,40 @@ def train_wav2vec2(
 
     criterion = nn.BCEWithLogitsLoss()
 
-    # Training loop
+    # Training loop with mixed precision + gradient accumulation
     best_val_auc = 0.0
     patience_counter = 0
     history = []
+    scaler_amp = torch.amp.GradScaler("cuda") if device.type == "cuda" else None
+    accum_steps = 4  # effective batch = batch_size * accum_steps
 
     for epoch in range(config.epochs):
         model.train()
         train_loss = 0.0
+        optimizer.zero_grad()
 
-        for batch_audio, batch_labels in train_loader:
+        for step, (batch_audio, batch_labels) in enumerate(train_loader):
             batch_audio = batch_audio.to(device)
             batch_labels = batch_labels.to(device)
 
-            optimizer.zero_grad()
-            logits, _ = model(batch_audio)
-            loss = criterion(logits.squeeze(-1), batch_labels)
-            loss.backward()
-            optimizer.step()
+            if scaler_amp is not None:
+                with torch.amp.autocast("cuda"):
+                    logits, _ = model(batch_audio)
+                    loss = criterion(logits.squeeze(-1), batch_labels) / accum_steps
+                scaler_amp.scale(loss).backward()
+                if (step + 1) % accum_steps == 0 or (step + 1) == len(train_loader):
+                    scaler_amp.step(optimizer)
+                    scaler_amp.update()
+                    optimizer.zero_grad()
+            else:
+                logits, _ = model(batch_audio)
+                loss = criterion(logits.squeeze(-1), batch_labels) / accum_steps
+                loss.backward()
+                if (step + 1) % accum_steps == 0 or (step + 1) == len(train_loader):
+                    optimizer.step()
+                    optimizer.zero_grad()
 
-            train_loss += loss.item()
+            train_loss += loss.item() * accum_steps
 
         avg_train_loss = train_loss / len(train_loader)
 
@@ -272,7 +286,11 @@ def train_wav2vec2(
         with torch.no_grad():
             for batch_audio, batch_labels in val_loader:
                 batch_audio = batch_audio.to(device)
-                logits, _ = model(batch_audio)
+                if scaler_amp is not None:
+                    with torch.amp.autocast("cuda"):
+                        logits, _ = model(batch_audio)
+                else:
+                    logits, _ = model(batch_audio)
                 probs = torch.sigmoid(logits.squeeze(-1))
                 val_probs.extend(probs.cpu().numpy())
                 val_labels.extend(batch_labels.numpy())
@@ -289,7 +307,8 @@ def train_wav2vec2(
             f"Epoch {epoch + 1}/{config.epochs} | "
             f"Loss: {avg_train_loss:.4f} | "
             f"Val Acc: {val_acc:.4f} | "
-            f"Val AUC: {val_auc:.4f}"
+            f"Val AUC: {val_auc:.4f}",
+            flush=True,
         )
 
         history.append({