Make NB4 fit 90 minute CPU submission limit

nb04_inference.py

@@ -1,17 +1,23 @@
 """
-BirdCLEF+ 2026 — Notebook 4 FAST INFERENCE
-
-Fix for timeout:
-  • Batch all 12 chunks of each soundscape per model instead of segment-by-segment.
-  • Use best B0 folds by default instead of all 5.
-  • No TTA by default.
-  • Uses CUDA if available, else CPU.
-  • Raw sigmoid outputs only; no threshold/calibration.
-
-Before run: set MODEL_DIR to Kaggle dataset containing b0_fold*.pt.
+BirdCLEF+ 2026 — Notebook 4 ULTRA-FAST CPU INFERENCE
+
+Designed for Kaggle submission CPU limit (~90 min).
+
+Speed choices:
+  • Uses ONLY best B0 fold by default: fold2.
+  • Computes predictions every 10 seconds by default (6 chunks/file), then duplicates
+    each prediction to fill adjacent 5-second rows.
+  • No TTA.
+  • Batched per soundscape.
+  • Raw sigmoid probabilities, no thresholds/calibration.
+
+If this finishes with time left, improve score by setting:
+  B0_FOLDS = [2, 4]
+  PREDICT_STRIDE_SEC = 5
+But for first valid CPU submission, keep defaults.
 """
 
-import os, gc, time
+import os, time, gc
 import numpy as np
 import pandas as pd
 import torch
@@ -28,25 +34,26 @@ COMP_DIR = "/kaggle/input/competitions/birdclef-2026"
 TEST_DIR = f"{COMP_DIR}/test_soundscapes"
 SAMPLE_SUB = f"{COMP_DIR}/sample_submission.csv"
 
-# CHANGE THIS to your Kaggle dataset path.
+# CHANGE to your Kaggle model dataset path.
 MODEL_DIR = "/kaggle/input/birdclef-b0-5fold"
-# If your .pt files are inside a models folder:
 # MODEL_DIR = "/kaggle/input/birdclef-b0-5fold/models"
 
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DEVICE = "cpu"          # CPU submission limit. Do not depend on GPU.
 SR = 32000
 DURATION = 5
 N_SAMPLES = SR * DURATION
 
-# FAST DEFAULT: use top validation folds first.
-# Fold scores: f2=0.9244, f4=0.8896, f0=0.8556, f1=0.8551, f3=0.8035
-B0_FOLDS = [2, 4, 0]        # fastest good ensemble. If still timeout, use [2]. If safe, use [2,4,0,1]
-USE_B3 = False
-B3_FOLDS = []
-USE_TTA = False             # keep False for Kaggle runtime
+# CPU-safe defaults
+B0_FOLDS = [2]           # best validation fold: 0.9244. Fastest valid submission.
+USE_TTA = False
+PREDICT_STRIDE_SEC = 10  # 10 = compute 6 chunks/file and duplicate to 12 rows. 5 = full 12 chunks.
 
-# Lower thread count can reduce CPU thrash in Kaggle containers.
-torch.set_num_threads(2)
+# CPU tuning
+try:
+    torch.set_num_threads(4)
+    torch.set_num_interop_threads(1)
+except Exception:
+    pass
 
 # =========================
 # LOAD SAMPLE
@@ -59,9 +66,9 @@ NUM_CLASSES = len(SPECIES)
 # MODEL
 # =========================
 class Model(nn.Module):
-    def __init__(self, backbone):
+    def __init__(self):
         super().__init__()
-        self.backbone = timm.create_model(backbone, pretrained=False, in_chans=3, features_only=True)
+        self.backbone = timm.create_model("tf_efficientnet_b0_ns", pretrained=False, in_chans=3, features_only=True)
         fi = self.backbone.feature_info
         ch = fi[-2]["num_chs"] + fi[-1]["num_chs"]
         self.pool = nn.AdaptiveAvgPool2d(1)
@@ -79,67 +86,57 @@ class Model(nn.Module):
 # =========================
 # LOAD MODELS
 # =========================
-def load_one(backbone, path):
-    m = Model(backbone)
-    state = torch.load(path, map_location="cpu")
-    m.load_state_dict(state, strict=False)
-    m.to(DEVICE)
-    m.eval()
-    return m
-
 MODELS = []
 for fold in B0_FOLDS:
     path = f"{MODEL_DIR}/b0_fold{fold}.pt"
     if os.path.exists(path):
-        MODELS.append(("b0", load_one("tf_efficientnet_b0_ns", path)))
+        m = Model()
+        state = torch.load(path, map_location="cpu")
+        m.load_state_dict(state, strict=False)
+        m.eval()
+        m.to(DEVICE)
+        MODELS.append(m)
         print("loaded:", path)
     else:
         print("missing:", path)
 
-if USE_B3:
-    for fold in B3_FOLDS:
-        path = f"{MODEL_DIR}/b3_fold{fold}.pt"
-        if os.path.exists(path):
-            MODELS.append(("b3", load_one("tf_efficientnet_b3_ns", path)))
-            print("loaded:", path)
-        else:
-            print("missing:", path)
-
 if len(MODELS) == 0:
-    raise ValueError(f"No models loaded. Check MODEL_DIR={MODEL_DIR}")
+    raise ValueError(f"No models loaded from MODEL_DIR={MODEL_DIR}. Check dataset path.")
 
-print("DEVICE:", DEVICE)
-print("ensemble size:", len(MODELS))
-print("B0_FOLDS:", B0_FOLDS, "USE_TTA:", USE_TTA)
+print("CPU ultra-fast config")
+print("models:", len(MODELS), "folds:", B0_FOLDS)
+print("PREDICT_STRIDE_SEC:", PREDICT_STRIDE_SEC)
 
 # =========================
-# FEATURES
+# FEATURE HELPERS
 # =========================
-def make_spec_np(chunk, n_fft, hop):
+def make_spec_np(chunk):
+    # Must match B0 training spec_a: n_fft=1024, hop=64, n_mels=128.
     mel = librosa.feature.melspectrogram(
-        y=chunk, sr=SR, n_fft=n_fft, hop_length=hop,
+        y=chunk, sr=SR, n_fft=1024, hop_length=64,
         n_mels=128, fmin=20, fmax=16000
     )
     mel = librosa.power_to_db(mel)
     mel = (mel - mel.min()) / (mel.max() - mel.min() + 1e-6)
-    return np.stack([mel] * 3).astype(np.float32)
-
-def get_chunks_12(wav):
-    chunks = []
-    for sec in range(0, 60, 5):
-        start = sec * SR
-        chunk = wav[start:start + N_SAMPLES]
-        if len(chunk) < N_SAMPLES:
-            chunk = np.pad(chunk, (0, N_SAMPLES - len(chunk)))
-        chunks.append(chunk.astype(np.float32))
-    return chunks
-
-def make_batch(chunks, model_name):
-    if model_name == "b0":
-        specs = [make_spec_np(c, 1024, 64) for c in chunks]
-    else:
-        specs = [make_spec_np(c, 2048, 512) for c in chunks]
-    return torch.tensor(np.stack(specs), dtype=torch.float32, device=DEVICE)
+    return np.stack([mel, mel, mel]).astype(np.float32)
+
+def chunk_at(wav, sec):
+    start = sec * SR
+    chunk = wav[start:start + N_SAMPLES]
+    if len(chunk) < N_SAMPLES:
+        chunk = np.pad(chunk, (0, N_SAMPLES - len(chunk)))
+    return chunk.astype(np.float32)
+
+def predict_chunks(chunks):
+    specs = [make_spec_np(c) for c in chunks]
+    x = torch.from_numpy(np.stack(specs)).to(DEVICE)
+    logits_sum = None
+    with torch.inference_mode():
+        for m in MODELS:
+            logits = m(x).detach().cpu().numpy()
+            logits_sum = logits if logits_sum is None else logits_sum + logits
+    logits = logits_sum / len(MODELS)
+    return (1.0 / (1.0 + np.exp(-logits))).astype(np.float32)
 
 # =========================
 # INFERENCE
@@ -147,7 +144,7 @@ def make_batch(chunks, model_name):
 files = sorted([f for f in os.listdir(TEST_DIR) if f.endswith((".ogg", ".wav", ".flac", ".mp3"))])
 print("test files:", len(files))
 
-all_rows = []
+all_row_ids = []
 all_preds = []
 t0 = time.time()
 
@@ -165,35 +162,41 @@ for file_idx, fname in enumerate(files):
         wav = wav.mean(axis=1)
     if sr != SR:
         wav = librosa.resample(wav, orig_sr=sr, target_sr=SR)
-
-    chunks = get_chunks_12(wav)
-    row_ids = [f"{stem}_{sec}" for sec in range(5, 65, 5)]
-
-    logits_sum = None
-    n_pred = 0
-
-    with torch.no_grad():
-        # Build b0 batch once if needed; b3 batch once if needed.
-        batch_cache = {}
-        for model_name, model in MODELS:
-            if model_name not in batch_cache:
-                batch_cache[model_name] = make_batch(chunks, model_name)
-            x = batch_cache[model_name]
-            logits = model(x).detach().float().cpu().numpy()  # shape (12, classes)
-            logits_sum = logits if logits_sum is None else logits_sum + logits
-            n_pred += 1
-            if DEVICE == "cuda":
-                torch.cuda.empty_cache()
-
-    avg_logits = logits_sum / max(n_pred, 1)
-    probs = 1.0 / (1.0 + np.exp(-avg_logits))
-
-    all_rows.extend(row_ids)
-    all_preds.append(probs.astype(np.float32))
-
-    if file_idx == 0 or (file_idx + 1) % 25 == 0:
-        elapsed = time.time() - t0
-        print(f"progress {file_idx+1}/{len(files)} elapsed={elapsed/60:.1f} min")
+    wav = wav.astype(np.float32)
+
+    # Standard row seconds: 5,10,...,60 with chunk starts 0,5,...,55.
+    if PREDICT_STRIDE_SEC <= 5:
+        start_secs = list(range(0, 60, 5))
+        chunks = [chunk_at(wav, s) for s in start_secs]
+        probs = predict_chunks(chunks)  # (12, C)
+        row_secs = list(range(5, 65, 5))
+        row_probs = probs
+    else:
+        # Compute every 10 sec: starts 0,10,20,30,40,50 = 6 predictions.
+        # Duplicate each prediction to adjacent 5-sec row.
+        start_secs = list(range(0, 60, PREDICT_STRIDE_SEC))
+        chunks = [chunk_at(wav, s) for s in start_secs]
+        probs6 = predict_chunks(chunks)  # (6, C)
+        row_secs = []
+        row_probs = []
+        for i, s in enumerate(start_secs):
+            # prediction for chunk s..s+5 fills row end s+5 and s+10
+            e1 = s + 5
+            e2 = s + 10
+            if e1 <= 60:
+                row_secs.append(e1)
+                row_probs.append(probs6[i])
+            if e2 <= 60:
+                row_secs.append(e2)
+                row_probs.append(probs6[i])
+        row_probs = np.stack(row_probs).astype(np.float32)
+
+    all_row_ids.extend([f"{stem}_{sec}" for sec in row_secs])
+    all_preds.append(row_probs)
+
+    if file_idx == 0 or (file_idx + 1) % 20 == 0:
+        elapsed = (time.time() - t0) / 60
+        print(f"progress {file_idx+1}/{len(files)} elapsed={elapsed:.1f} min")
 
     gc.collect()
 
@@ -201,14 +204,14 @@ for file_idx, fname in enumerate(files):
 # SUBMISSION
 # =========================
 if len(all_preds) == 0:
-    pred_arr = np.zeros((len(all_rows), NUM_CLASSES), dtype=np.float32)
+    pred_arr = np.zeros((len(all_row_ids), NUM_CLASSES), dtype=np.float32)
 else:
     pred_arr = np.vstack(all_preds)
 
 sub = pd.DataFrame(pred_arr, columns=SPECIES)
-sub.insert(0, "row_id", all_rows)
+sub.insert(0, "row_id", all_row_ids)
 
-# Align exactly with sample submission
+# Align exactly with sample_submission. Missing rows filled 0, but should not be missing.
 sub = sample[["row_id"]].merge(sub, on="row_id", how="left").fillna(0)
 sub = sub[sample.columns]
 
@@ -219,8 +222,9 @@ sub.to_csv("submission.csv", index=False)
 
 print("SUBMISSION READY")
 print("shape:", sub.shape)
-print("models:", len(MODELS), "folds:", B0_FOLDS, "DEVICE:", DEVICE)
+print("models:", len(MODELS), "folds:", B0_FOLDS)
+print("stride:", PREDICT_STRIDE_SEC)
 print("mean prob:", float(sub[SPECIES].values.mean()))
 print("max prob:", float(sub[SPECIES].values.max()))
 print("nonzero ratio:", float((sub[SPECIES].values > 0).mean()))
-print("total elapsed min:", (time.time() - t0) / 60)
+print("elapsed min:", (time.time() - t0) / 60)