Upload finetune_large_geolifeclef.py with huggingface_hub

finetune_large_geolifeclef.py

@@ -0,0 +1,264 @@
+"""
+Finetune BFM Large model on GeoLifeCLEF 500 species data.
+
+Based on the paper's pipeline (arXiv:2507.09080v2):
+  1. Load BFM Large pretrained model from safetensors
+  2. Wrap with BFMRaw (replace encoder/decoder with species-specific ones)
+  3. Train with L1 loss + AdamW + CosineAnnealing on GeoLifeCLEF data
+  4. Save finetuned checkpoint
+
+Usage:
+    conda run -n bfm python finetune_large_geolifeclef.py 2>&1 | tee finetune_large.log
+"""
+
+import math
+import os
+import sys
+import time
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+# Add bfm-model to path
+PROJECT_ROOT = Path(__file__).resolve().parent
+BFM_MODEL_DIR = PROJECT_ROOT / "bfm-model"
+sys.path.insert(0, str(BFM_MODEL_DIR))
+
+SAFETENSORS_PATH = PROJECT_ROOT / "bfm-pretrained" / "bfm-pretrain-large.safetensors"
+OUTPUT_DIR = PROJECT_ROOT / "outputs_finetune_large"
+CHECKPOINT_DIR = OUTPUT_DIR / "checkpoints"
+
+# ─── Training hyperparameters (from finetune_config.yaml) ───
+NUM_SPECIES = 500
+BATCH_SIZE = 1
+NUM_EPOCHS = 100
+LEARNING_RATE = 3e-4
+VAL_EVERY = 5
+NUM_WORKERS = 8
+
+# ─── Model config (Large) ───
+MODEL_CONFIG = {
+    "embed_dim": 512, "depth": 10, "patch_size": 8,
+    "swin_backbone_size": "large", "num_heads": 16, "head_dim": 64,
+    "H": 160, "W": 280, "num_latent_tokens": 8,
+    "perceiver_latents": 16100, "T": 2,
+}
+
+SWIN_LARGE_CONFIG = {
+    "swin_encoder_depths": (2, 2, 2), "swin_encoder_num_heads": (8, 16, 32),
+    "swin_decoder_depths": (2, 2, 2), "swin_decoder_num_heads": (32, 16, 8),
+    "swin_window_size": (1, 4, 5), "swin_mlp_ratio": 4.0,
+    "swin_qkv_bias": True, "swin_drop_rate": 0.0,
+    "swin_attn_drop_rate": 0.0, "swin_drop_path_rate": 0.1,
+    "use_lora": False,
+}
+
+SPECIES_VARS = [
+    "1340361", "1340503", "1536449", "1898286", "1920506", "2430567",
+    "2431885", "2433433", "2434779", "2435240", "2435261", "2437394",
+    "2441454", "2473958", "2491534", "2891770", "3034825", "4408498",
+    "5218786", "5219073", "5219173", "5219219", "5844449", "8002952",
+    "8077224", "8894817", "8909809", "9809229",
+]
+SURFACE_VARS = ["t2m", "msl", "slt", "z", "u10", "v10", "lsm"]
+EDAPHIC_VARS = ["swvl1", "swvl2", "stl1", "stl2"]
+ATMOS_VARS = ["z", "t", "u", "v", "q"]
+CLIMATE_VARS = [
+    "smlt", "tp", "csfr", "avg_sdswrf", "avg_snswrf", "avg_snlwrf",
+    "avg_tprate", "avg_sdswrfcs", "sd", "t2m", "d2m",
+]
+VEGETATION_VARS = ["NDVI"]
+LAND_VARS = ["Land"]
+AGRICULTURE_VARS = ["Agriculture", "Arable", "Cropland"]
+FOREST_VARS = ["Forest"]
+REDLIST_VARS = ["RLI"]
+MISC_VARS = ["avg_slhtf", "avg_pevr"]
+ATMOS_LEVELS = [1000, 925, 850, 700, 600, 500, 400, 300, 250, 200, 150, 100, 50]
+
+
+def build_base_model():
+    """Build BFM Large model."""
+    from bfm_model.bfm.model import BFM
+    model = BFM(
+        surface_vars=SURFACE_VARS, edaphic_vars=EDAPHIC_VARS,
+        atmos_vars=ATMOS_VARS, climate_vars=CLIMATE_VARS,
+        species_vars=SPECIES_VARS, vegetation_vars=VEGETATION_VARS,
+        land_vars=LAND_VARS, agriculture_vars=AGRICULTURE_VARS,
+        forest_vars=FOREST_VARS, redlist_vars=REDLIST_VARS,
+        misc_vars=MISC_VARS, atmos_levels=ATMOS_LEVELS,
+        species_num=len(SPECIES_VARS),
+        H=MODEL_CONFIG["H"], W=MODEL_CONFIG["W"],
+        num_latent_tokens=MODEL_CONFIG["num_latent_tokens"],
+        backbone_type="swin", patch_size=MODEL_CONFIG["patch_size"],
+        embed_dim=MODEL_CONFIG["embed_dim"],
+        num_heads=MODEL_CONFIG["num_heads"],
+        head_dim=MODEL_CONFIG["head_dim"],
+        depth=MODEL_CONFIG["depth"],
+        perceiver_latents=MODEL_CONFIG["perceiver_latents"],
+        batch_size=1, td_learning=True, use_mask="no",
+        **SWIN_LARGE_CONFIG,
+    )
+    return model
+
+
+def load_pretrained_weights(model):
+    """Load pretrained safetensors weights into base model."""
+    from safetensors.torch import load_file
+    print(f"Loading weights from {SAFETENSORS_PATH.name}...")
+    state = load_file(str(SAFETENSORS_PATH), device="cpu")
+    missing, unexpected = model.load_state_dict(state, strict=False)
+    alias_missing = [k for k in missing if "_latent_parameter_list" in k]
+    real_missing = [k for k in missing if "_latent_parameter_list" not in k]
+    print(f"  Total missing: {len(missing)} ({len(alias_missing)} aliases)")
+    print(f"  Real missing: {len(real_missing)}")
+    print(f"  Unexpected: {len(unexpected)}")
+    return model
+
+
+def train_epoch(model, dataloader, optimizer, criterion, scheduler, device):
+    """One training epoch."""
+    model.train()
+    epoch_loss = 0.0
+    for batch_idx, sample in enumerate(dataloader):
+        batch = sample["batch"]
+        batch["species_distribution"] = batch["species_distribution"].to(device)
+        targets = sample["target"].to(device)
+        optimizer.zero_grad()
+        outputs = model(batch)
+        loss = criterion(outputs, targets)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        optimizer.step()
+        scheduler.step()
+        epoch_loss += loss.item()
+        if (batch_idx + 1) % 1 == 0:
+            print(f"    Batch {batch_idx+1}/{len(dataloader)}, Loss: {loss.item():.6f}")
+    return epoch_loss / max(len(dataloader), 1)
+
+
+def validate_epoch(model, dataloader, criterion, device):
+    """One validation epoch."""
+    model.eval()
+    epoch_loss = 0.0
+    with torch.inference_mode():
+        for sample in dataloader:
+            batch = sample["batch"]
+            batch["species_distribution"] = batch["species_distribution"].to(device)
+            targets = sample["target"].to(device)
+            outputs = model(batch)
+            loss = criterion(outputs, targets)
+            epoch_loss += loss.item()
+    return epoch_loss / max(len(dataloader), 1)
+
+
+def save_checkpoint(model, optimizer, epoch, loss, path):
+    """Save training checkpoint."""
+    os.makedirs(path, exist_ok=True)
+    filepath = path / "best_checkpoint.pth"
+    torch.save({
+        "epoch": epoch,
+        "model_state_dict": model.state_dict(),
+        "optimizer_state_dict": optimizer.state_dict(),
+        "loss": loss,
+    }, filepath)
+    print(f"  Checkpoint saved: epoch={epoch}, loss={loss:.6f}")
+
+
+def main():
+    print("=" * 70)
+    print("BFM Large Model GeoLifeCLEF Finetuning")
+    print("=" * 70)
+
+    assert SAFETENSORS_PATH.exists(), f"Weights not found: {SAFETENSORS_PATH}"
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+    torch.set_float32_matmul_precision("highest")
+
+    # 1. Build base model and load pretrained weights
+    print("\nBuilding BFM Large model...")
+    base_model = build_base_model()
+    total_params = sum(p.numel() for p in base_model.parameters())
+    print(f"Base model parameters: {total_params / 1e6:.1f}M")
+
+    base_model = load_pretrained_weights(base_model)
+
+    # 2. Wrap with BFMRaw for species finetuning
+    print("\nWrapping with BFMRaw for species finetuning...")
+    from bfm_finetune.bfm_mod import BFMRaw
+    model = BFMRaw(base_model=base_model, n_species=NUM_SPECIES, mode="train")
+    model.to(device)
+
+    # 3. Setup datasets
+    print("\nLoading GeoLifeCLEF datasets...")
+    from bfm_finetune.dataloaders.geolifeclef_species.dataloader import GeoLifeCLEFSpeciesDataset
+    from bfm_finetune.dataloaders.dataloader_utils import custom_collate_fn
+
+    train_dataset = GeoLifeCLEFSpeciesDataset(
+        num_species=NUM_SPECIES, mode="train", negative_lon_mode="ignore",
+    )
+    val_dataset = GeoLifeCLEFSpeciesDataset(
+        num_species=NUM_SPECIES, mode="val", negative_lon_mode="ignore",
+    )
+
+    train_dataloader = DataLoader(
+        train_dataset, batch_size=BATCH_SIZE, shuffle=True,
+        collate_fn=custom_collate_fn, num_workers=NUM_WORKERS,
+    )
+    val_dataloader = DataLoader(
+        val_dataset, batch_size=1, shuffle=False,
+        collate_fn=custom_collate_fn, num_workers=NUM_WORKERS,
+    )
+    print(f"Train: {len(train_dataset)} samples, Val: {len(val_dataset)} samples")
+
+    # 4. Setup optimizer, scheduler, loss
+    optimizer = torch.optim.AdamW(
+        model.parameters(), lr=LEARNING_RATE,
+        weight_decay=0.0001, betas=(0.9, 0.95), eps=1e-8,
+    )
+    criterion = nn.L1Loss()
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+        optimizer, T_max=NUM_EPOCHS * len(train_dataloader),
+        eta_min=LEARNING_RATE / 10,
+    )
+
+    # 5. Training loop
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+    best_loss = float("inf")
+    start_time = time.time()
+
+    print(f"\nStarting training: {NUM_EPOCHS} epochs, batch_size={BATCH_SIZE}")
+    print("-" * 70)
+
+    for epoch in range(1, NUM_EPOCHS + 1):
+        t0 = time.time()
+        train_loss = train_epoch(model, train_dataloader, optimizer, criterion, scheduler, device)
+        epoch_time = time.time() - t0
+
+        log_msg = f"Epoch {epoch}/{NUM_EPOCHS}, Train Loss: {train_loss:.6f}, Time: {epoch_time:.1f}s"
+
+        if epoch % VAL_EVERY == 0:
+            val_loss = validate_epoch(model, val_dataloader, criterion, device)
+            log_msg += f", Val Loss: {val_loss:.6f}"
+
+            if val_loss < best_loss:
+                best_loss = val_loss
+                save_checkpoint(model, optimizer, epoch, best_loss, CHECKPOINT_DIR)
+        else:
+            # Save based on train loss if no validation this epoch
+            if train_loss < best_loss:
+                best_loss = train_loss
+                save_checkpoint(model, optimizer, epoch, best_loss, CHECKPOINT_DIR)
+
+        print(log_msg)
+
+    total_time = time.time() - start_time
+    print(f"\nTraining complete! Total time: {total_time/60:.1f} minutes")
+    print(f"Best loss: {best_loss:.6f}")
+    print(f"Checkpoint saved to: {CHECKPOINT_DIR}")
+
+
+if __name__ == "__main__":
+    main()