Add ML diagnostics tooling (W&B, TensorBoard, Grad-CAM, profiler) and gin configs
- Add unified experiment tracker (W&B + TensorBoard) with graceful fallback - Add gradient norm monitoring per param group (MONA, LoRA, MLP, gates, tau) - Add Grad-CAM visualization for DINOv3 drone/satellite encoders - Add PyTorch Profiler wrapper + torchinfo model summary - Add gin-config support to train_gtauav.py with CLI overrides - Add v3 gin configs: gtauav_balanced, gtauav_baseline, gtauav_text_heavy, gtauav_image_heavy - Generate metric plots every epoch (not just on eval) - Set default epochs to 10 - Update README and CLAUDE.md with new tooling and usage docs Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
@@ -4,6 +4,9 @@ from __future__ import annotations
|
||||
|
||||
Asymmetric DINOv3 encoders (drone LVD + satellite SAT) with LRSCLIP text fusion.
|
||||
Single InfoNCE loss: query(drone+text) vs gallery(satellite).
|
||||
|
||||
Supports gin-config, W&B, TensorBoard, Grad-CAM, gradient monitoring,
|
||||
PyTorch Profiler, and torchinfo model summary.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
@@ -16,6 +19,7 @@ from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
import coloredlogs
|
||||
import gin
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
@@ -28,6 +32,9 @@ from tqdm import tqdm
|
||||
from src.datasets.gtauav_dataset import GTAUAVDataset, collate_gtauav_batch
|
||||
from src.losses.multi_infonce import InfoNCELoss
|
||||
from src.training.plot_metrics import generate_plots
|
||||
from src.training.trackers import ExperimentTracker
|
||||
from src.training.grad_monitor import compute_gradient_norms, log_gradient_summary
|
||||
from src.training.profiling import TrainingProfiler, print_model_summary
|
||||
from src.models.asymmetric_encoder import (
|
||||
AsymmetricEncoder,
|
||||
get_dino_transform,
|
||||
@@ -48,6 +55,7 @@ _DINO_SAT = "nn_models/DINO_SAT/model.safetensors"
|
||||
_LRSCLIP = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt"
|
||||
|
||||
|
||||
@gin.configurable
|
||||
@dataclass
|
||||
class TrainConfigGTAUAV:
|
||||
"""Training configuration for GTA-UAV experiment."""
|
||||
@@ -69,7 +77,7 @@ class TrainConfigGTAUAV:
|
||||
# Training.
|
||||
resume_from: str | None = None # path to checkpoint for resuming
|
||||
output_dir: str = "out/gtauav/with_text"
|
||||
epochs: int = 20
|
||||
epochs: int = 10
|
||||
batch_size: int = 8
|
||||
num_workers: int = 4
|
||||
learning_rate: float = 1e-4
|
||||
@@ -89,6 +97,20 @@ class TrainConfigGTAUAV:
|
||||
weight_g2q: float = 0.4
|
||||
learnable_temperature: bool = True
|
||||
|
||||
# Tracking & diagnostics.
|
||||
use_wandb: bool = False
|
||||
use_tb: bool = True
|
||||
wandb_project: str = "caption-test-gtauav"
|
||||
wandb_run_name: str | None = None
|
||||
wandb_entity: str | None = None
|
||||
log_grad_norms: bool = True
|
||||
use_gradcam: bool = False
|
||||
gradcam_every: int = 5 # Grad-CAM every N epochs
|
||||
gradcam_samples: int = 8
|
||||
use_profiler: bool = False
|
||||
profiler_warmup: int = 3
|
||||
profiler_active: int = 5
|
||||
|
||||
|
||||
def _set_seed(seed: int) -> None:
|
||||
import random as _random
|
||||
@@ -157,7 +179,7 @@ def _evaluate(
|
||||
all_query: list[torch.Tensor] = []
|
||||
all_gallery: list[torch.Tensor] = []
|
||||
|
||||
for batch in tqdm(loader, desc=" 🔎 Evaluating", unit="batch", leave=False):
|
||||
for batch in tqdm(loader, desc=" eval", unit="batch", leave=False):
|
||||
drone_img = batch["drone_img"].to(device, non_blocking=True)
|
||||
sat_img = batch["sat_img"].to(device, non_blocking=True)
|
||||
|
||||
@@ -240,7 +262,7 @@ def _clear_vram() -> None:
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
allocated = torch.cuda.memory_allocated() / 1e9
|
||||
LOGGER.info("🧹 VRAM cleared. Current usage: %.2f GB", allocated)
|
||||
LOGGER.info("VRAM cleared. Current usage: %.2f GB", allocated)
|
||||
|
||||
|
||||
def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
@@ -259,12 +281,23 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
with (output_dir / "config.json").open("w") as f:
|
||||
json.dump(vars(cfg), f, indent=2)
|
||||
|
||||
# --- Experiment tracker (W&B + TensorBoard) ---
|
||||
tracker = ExperimentTracker(
|
||||
output_dir=output_dir,
|
||||
config=vars(cfg),
|
||||
use_wandb=cfg.use_wandb,
|
||||
use_tb=cfg.use_tb,
|
||||
wandb_project=cfg.wandb_project,
|
||||
wandb_run_name=cfg.wandb_run_name,
|
||||
wandb_entity=cfg.wandb_entity,
|
||||
)
|
||||
|
||||
# Model.
|
||||
start_epoch = 0
|
||||
resume_ckpt = None
|
||||
|
||||
if cfg.resume_from is not None:
|
||||
LOGGER.info("🔄 Resuming from %s", cfg.resume_from)
|
||||
LOGGER.info("Resuming from %s", cfg.resume_from)
|
||||
model, resume_ckpt = AsymmetricEncoder.load_checkpoint(
|
||||
cfg.resume_from,
|
||||
dino_web_path=cfg.dino_web_path,
|
||||
@@ -274,8 +307,8 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
)
|
||||
start_epoch = resume_ckpt.get("epoch", -1) + 1
|
||||
else:
|
||||
mode_str = "🚫 baseline (no text)" if cfg.baseline_mode else "📝 with text (L1/L2/L3)"
|
||||
LOGGER.info("🏗️ Building model — %s", mode_str)
|
||||
mode_str = "baseline (no text)" if cfg.baseline_mode else "with text (L1/L2/L3)"
|
||||
LOGGER.info("Building model — %s", mode_str)
|
||||
model = AsymmetricEncoder(
|
||||
dino_web_path=cfg.dino_web_path,
|
||||
dino_sat_path=cfg.dino_sat_path,
|
||||
@@ -288,10 +321,18 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
n_trainable = sum(p.numel() for p in model.trainable_parameters())
|
||||
n_total = sum(p.numel() for p in model.parameters())
|
||||
LOGGER.info(
|
||||
"🧮 trainable=%s (%.2f%%) total=%s",
|
||||
"trainable=%s (%.2f%%) total=%s",
|
||||
f"{n_trainable:,}", 100.0 * n_trainable / max(n_total, 1), f"{n_total:,}",
|
||||
)
|
||||
|
||||
# --- Model summary (torchinfo) ---
|
||||
model_summary = print_model_summary(model, device=cfg.device)
|
||||
(output_dir / "model_summary.txt").write_text(model_summary)
|
||||
|
||||
# --- W&B model watching (gradient + weight histograms) ---
|
||||
if tracker.has_wandb:
|
||||
tracker.watch_model(model, log_freq=50)
|
||||
|
||||
# Loss.
|
||||
loss_fn = InfoNCELoss(
|
||||
temperature_init=cfg.tau_init,
|
||||
@@ -301,7 +342,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
learnable_temperature=cfg.learnable_temperature,
|
||||
)
|
||||
LOGGER.info(
|
||||
"🌡️ Temperature: %s (init=%.3f)",
|
||||
"Temperature: %s (init=%.3f)",
|
||||
"learnable" if cfg.learnable_temperature else "cosine schedule",
|
||||
cfg.tau_init,
|
||||
)
|
||||
@@ -345,7 +386,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
pin_memory=True,
|
||||
)
|
||||
|
||||
LOGGER.info("📦 train=%d test=%d batch=%d", len(train_ds), len(test_ds), cfg.batch_size)
|
||||
LOGGER.info("train=%d test=%d batch=%d", len(train_ds), len(test_ds), cfg.batch_size)
|
||||
|
||||
# Optimizer — per-group LR (text encoder gets lower LR).
|
||||
param_groups = _build_param_groups(model, cfg.learning_rate, cfg.text_lr_factor)
|
||||
@@ -358,7 +399,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
lr_info = f"proj={cfg.learning_rate:.0e}"
|
||||
if not cfg.baseline_mode:
|
||||
lr_info += f" text={cfg.learning_rate * cfg.text_lr_factor:.0e}"
|
||||
LOGGER.info("⚙️ Optimizer: AdamW LR: %s warmup=%d epochs", lr_info, cfg.warmup_epochs)
|
||||
LOGGER.info("Optimizer: AdamW LR: %s warmup=%d epochs", lr_info, cfg.warmup_epochs)
|
||||
|
||||
# Scheduler — cosine with linear warmup.
|
||||
steps_per_epoch = len(train_loader)
|
||||
@@ -377,18 +418,31 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
if resume_ckpt is not None:
|
||||
if "optimizer_state" in resume_ckpt:
|
||||
optimizer.load_state_dict(resume_ckpt["optimizer_state"])
|
||||
LOGGER.info("🔄 Optimizer state restored")
|
||||
LOGGER.info("Optimizer state restored")
|
||||
if "loss_state" in resume_ckpt:
|
||||
loss_fn.load_state_dict(resume_ckpt["loss_state"])
|
||||
LOGGER.info("🔄 Loss state restored (tau=%.4f)", loss_fn.current_temperature)
|
||||
LOGGER.info("Loss state restored (tau=%.4f)", loss_fn.current_temperature)
|
||||
# Set scheduler last_epoch so it resumes at the correct LR.
|
||||
scheduler.last_epoch = start_epoch * steps_per_epoch
|
||||
LOGGER.info("🔄 Resuming from epoch %d", start_epoch)
|
||||
LOGGER.info("Resuming from epoch %d", start_epoch)
|
||||
|
||||
history: list[dict] = []
|
||||
csv_logger = CSVLogger(output_dir)
|
||||
|
||||
LOGGER.info("🚀 Starting training for %d epochs (from epoch %d)", cfg.epochs, start_epoch)
|
||||
# --- Optional profiler (first epoch only) ---
|
||||
profiler = None
|
||||
if cfg.use_profiler and start_epoch == 0:
|
||||
profiler = TrainingProfiler(
|
||||
output_dir=output_dir,
|
||||
n_warmup=cfg.profiler_warmup,
|
||||
n_active=cfg.profiler_active,
|
||||
)
|
||||
profiler.start()
|
||||
|
||||
LOGGER.info("Starting training for %d epochs (from epoch %d)", cfg.epochs, start_epoch)
|
||||
|
||||
global_step = start_epoch * steps_per_epoch
|
||||
best_r1 = 0.0
|
||||
|
||||
for epoch in range(start_epoch, cfg.epochs):
|
||||
model.train()
|
||||
@@ -398,7 +452,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
|
||||
pbar = tqdm(
|
||||
train_loader,
|
||||
desc=f" 🏋️ Epoch {epoch + 1}/{cfg.epochs}",
|
||||
desc=f" Epoch {epoch + 1}/{cfg.epochs}",
|
||||
unit="batch",
|
||||
leave=False,
|
||||
)
|
||||
@@ -439,15 +493,34 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
model.trainable_parameters(),
|
||||
max_norm=cfg.grad_clip,
|
||||
)
|
||||
|
||||
# --- Gradient monitoring (after unscale, before step) ---
|
||||
if cfg.log_grad_norms and n_batches % 50 == 0:
|
||||
grad_norms = compute_gradient_norms(model, loss_fn)
|
||||
tracker.log_gradients(epoch, grad_norms, step=global_step)
|
||||
if n_batches == 0:
|
||||
log_gradient_summary(grad_norms)
|
||||
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
with warnings.catch_warnings():
|
||||
warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*")
|
||||
scheduler.step()
|
||||
|
||||
# --- Per-step tracking ---
|
||||
step_metrics = {
|
||||
"loss": float(total_loss.item()),
|
||||
"temperature": float(loss_dict["temperature"].item()),
|
||||
"gate_q": float(loss_dict["gate_q"].item()),
|
||||
"gate_g": float(loss_dict["gate_g"].item()),
|
||||
"lr": optimizer.param_groups[0]["lr"],
|
||||
}
|
||||
tracker.log_train(epoch, step_metrics, step=global_step)
|
||||
|
||||
for key, val in loss_dict.items():
|
||||
agg[key] = agg.get(key, 0.0) + float(val.item())
|
||||
n_batches += 1
|
||||
global_step += 1
|
||||
|
||||
pbar.set_postfix(
|
||||
loss=f"{total_loss.item():.3f}",
|
||||
@@ -456,11 +529,18 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
gg=f"{loss_dict['gate_g'].item():.3f}",
|
||||
)
|
||||
|
||||
# --- Profiler step ---
|
||||
if profiler is not None:
|
||||
profiler.step()
|
||||
if profiler.is_done(n_batches):
|
||||
profiler.export()
|
||||
profiler = None
|
||||
|
||||
elapsed = time.time() - epoch_start
|
||||
|
||||
means = {k: v / max(n_batches, 1) for k, v in agg.items()}
|
||||
LOGGER.info(
|
||||
"📈 epoch=%d time=%.1fs lr=%.2e loss=%.4f tau=%.4f gate_q=%.4f gate_g=%.4f",
|
||||
"epoch=%d time=%.1fs lr=%.2e loss=%.4f tau=%.4f gate_q=%.4f gate_g=%.4f",
|
||||
epoch, elapsed,
|
||||
optimizer.param_groups[0]["lr"],
|
||||
means.get("total", 0.0),
|
||||
@@ -475,8 +555,14 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
"train": means,
|
||||
}
|
||||
|
||||
# Log train metrics to CSV.
|
||||
# Log train metrics to CSV + generate plots every epoch.
|
||||
csv_logger.log_train(epoch, means, optimizer.param_groups[0]["lr"], elapsed)
|
||||
generate_plots(csv_logger.log_dir)
|
||||
|
||||
# --- Log VRAM usage ---
|
||||
if torch.cuda.is_available():
|
||||
vram_gb = torch.cuda.max_memory_allocated() / 1e9
|
||||
tracker.log_scalar("system/vram_peak_gb", vram_gb, step=global_step)
|
||||
|
||||
# Evaluation.
|
||||
if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1:
|
||||
@@ -484,8 +570,16 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
epoch_record["val"] = val_metrics
|
||||
csv_logger.log_val(epoch, val_metrics)
|
||||
generate_plots(csv_logger.log_dir)
|
||||
tracker.log_val(epoch, val_metrics, step=global_step)
|
||||
|
||||
# Track best R@1.
|
||||
r1 = val_metrics.get("r@1_q2g", 0.0)
|
||||
if r1 > best_r1:
|
||||
best_r1 = r1
|
||||
tracker.log_scalar("val/best_r@1_q2g", best_r1, step=global_step)
|
||||
|
||||
LOGGER.info(
|
||||
"🎯 val epoch=%d R@1=%.4f R@5=%.4f R@10=%.4f gate_q=%.4f gate_g=%.4f",
|
||||
"val epoch=%d R@1=%.4f R@5=%.4f R@10=%.4f gate_q=%.4f gate_g=%.4f",
|
||||
epoch,
|
||||
val_metrics.get("r@1_q2g", 0.0),
|
||||
val_metrics.get("r@5_q2g", 0.0),
|
||||
@@ -494,6 +588,27 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
val_metrics.get("gate_g", 1.0),
|
||||
)
|
||||
|
||||
# --- Grad-CAM visualization ---
|
||||
if cfg.use_gradcam and (epoch + 1) % cfg.gradcam_every == 0:
|
||||
from src.training.gradcam import generate_gradcam_samples
|
||||
overlays = generate_gradcam_samples(
|
||||
model=model,
|
||||
dataloader=test_loader,
|
||||
device=cfg.device,
|
||||
output_dir=str(output_dir),
|
||||
n_samples=cfg.gradcam_samples,
|
||||
epoch=epoch,
|
||||
)
|
||||
# Log first few overlays to tracker.
|
||||
for i, overlay in enumerate(overlays[:4]):
|
||||
kind = "drone" if i % 2 == 0 else "sat"
|
||||
tracker.log_image(
|
||||
f"gradcam/{kind}_{i//2}",
|
||||
overlay,
|
||||
step=global_step,
|
||||
caption=f"Epoch {epoch} {kind} Grad-CAM",
|
||||
)
|
||||
|
||||
history.append(epoch_record)
|
||||
|
||||
# Save checkpoint.
|
||||
@@ -506,7 +621,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
},
|
||||
path=output_dir / f"ckpt_epoch{epoch:03d}.pt",
|
||||
)
|
||||
LOGGER.info("💾 Checkpoint saved: ckpt_epoch%03d.pt", epoch)
|
||||
LOGGER.info("Checkpoint saved: ckpt_epoch%03d.pt", epoch)
|
||||
|
||||
# Save history.
|
||||
history_path = output_dir / "history.json"
|
||||
@@ -514,7 +629,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
json.dump(history, f, indent=2)
|
||||
|
||||
# Save final eval report.
|
||||
LOGGER.info("🔎 Running final evaluation...")
|
||||
LOGGER.info("Running final evaluation...")
|
||||
final_metrics = _evaluate(model, test_loader, cfg.device)
|
||||
report = {
|
||||
"config": vars(cfg),
|
||||
@@ -525,9 +640,25 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
with report_path.open("w", encoding="utf-8") as f:
|
||||
json.dump(report, f, indent=2)
|
||||
|
||||
LOGGER.info("✅ Training complete. Report: %s", report_path)
|
||||
# --- Log final summary to W&B ---
|
||||
tracker.log_summary({
|
||||
"best_r@1_q2g": best_r1,
|
||||
"final_r@1_q2g": final_metrics.get("r@1_q2g", 0.0),
|
||||
"final_r@5_q2g": final_metrics.get("r@5_q2g", 0.0),
|
||||
"final_r@10_q2g": final_metrics.get("r@10_q2g", 0.0),
|
||||
"final_gate_q": final_metrics.get("gate_q", 1.0),
|
||||
"final_gate_g": final_metrics.get("gate_g", 1.0),
|
||||
})
|
||||
|
||||
# --- Cleanup profiler if still running ---
|
||||
if profiler is not None:
|
||||
profiler.export()
|
||||
|
||||
tracker.close()
|
||||
|
||||
LOGGER.info("Training complete. Report: %s", report_path)
|
||||
LOGGER.info(
|
||||
"📊 Final — R@1=%.4f R@5=%.4f R@10=%.4f gate_q=%.4f gate_g=%.4f",
|
||||
"Final — R@1=%.4f R@5=%.4f R@10=%.4f gate_q=%.4f gate_g=%.4f",
|
||||
final_metrics.get("r@1_q2g", 0.0),
|
||||
final_metrics.get("r@5_q2g", 0.0),
|
||||
final_metrics.get("r@10_q2g", 0.0),
|
||||
@@ -538,6 +669,10 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="GTA-UAV caption test training.")
|
||||
parser.add_argument(
|
||||
"--config", type=str, default=None,
|
||||
help="Path to gin config file (e.g. conf/gtauav_balanced.gin).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--baseline", action="store_true",
|
||||
help="Run baseline mode (no text).",
|
||||
@@ -555,47 +690,86 @@ def main() -> None:
|
||||
help="Path to seg_filter.json for excluding bad images.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch-size", type=int, default=8,
|
||||
"--batch-size", type=int, default=None,
|
||||
help="Batch size.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--epochs", type=int, default=20,
|
||||
"--epochs", type=int, default=None,
|
||||
help="Number of epochs.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lr", type=float, default=1e-4,
|
||||
"--lr", type=float, default=None,
|
||||
help="Learning rate for projections.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text-lr-factor", type=float, default=0.1,
|
||||
"--text-lr-factor", type=float, default=None,
|
||||
help="Text encoder LR = lr * factor (default 0.1 = 10x lower).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup-epochs", type=int, default=2,
|
||||
"--warmup-epochs", type=int, default=None,
|
||||
help="Linear warmup epochs.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--init-gate", type=float, default=0.7,
|
||||
"--init-gate", type=float, default=None,
|
||||
help="Initial gate value (image weight).",
|
||||
)
|
||||
# Tracking flags.
|
||||
parser.add_argument("--wandb", action="store_true", help="Enable W&B tracking.")
|
||||
parser.add_argument("--no-tb", action="store_true", help="Disable TensorBoard.")
|
||||
parser.add_argument("--gradcam", action="store_true", help="Enable Grad-CAM visualization.")
|
||||
parser.add_argument("--profile", action="store_true", help="Enable PyTorch profiler (first epoch).")
|
||||
parser.add_argument("--no-grad-norms", action="store_true", help="Disable gradient norm logging.")
|
||||
# Gin overrides.
|
||||
parser.add_argument(
|
||||
"--gin-param", type=str, nargs="*", default=[],
|
||||
help="Gin parameter overrides (e.g. 'TrainConfigGTAUAV.epochs=30').",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
cfg = TrainConfigGTAUAV()
|
||||
cfg.baseline_mode = args.baseline
|
||||
cfg.resume_from = args.resume
|
||||
cfg.batch_size = args.batch_size
|
||||
cfg.epochs = args.epochs
|
||||
cfg.learning_rate = args.lr
|
||||
cfg.text_lr_factor = args.text_lr_factor
|
||||
cfg.warmup_epochs = args.warmup_epochs
|
||||
cfg.init_gate = args.init_gate
|
||||
# Parse gin config if provided.
|
||||
if args.config is not None:
|
||||
gin.parse_config_file(args.config)
|
||||
if args.gin_param:
|
||||
gin.parse_config(args.gin_param)
|
||||
|
||||
# Create config (gin bindings apply via @gin.configurable).
|
||||
cfg = TrainConfigGTAUAV()
|
||||
|
||||
# CLI overrides take priority over gin.
|
||||
if args.baseline:
|
||||
cfg.baseline_mode = True
|
||||
if args.resume is not None:
|
||||
cfg.resume_from = args.resume
|
||||
if args.batch_size is not None:
|
||||
cfg.batch_size = args.batch_size
|
||||
if args.epochs is not None:
|
||||
cfg.epochs = args.epochs
|
||||
if args.lr is not None:
|
||||
cfg.learning_rate = args.lr
|
||||
if args.text_lr_factor is not None:
|
||||
cfg.text_lr_factor = args.text_lr_factor
|
||||
if args.warmup_epochs is not None:
|
||||
cfg.warmup_epochs = args.warmup_epochs
|
||||
if args.init_gate is not None:
|
||||
cfg.init_gate = args.init_gate
|
||||
if args.filter_meta is not None:
|
||||
cfg.filter_meta = args.filter_meta
|
||||
|
||||
# Tracking overrides.
|
||||
if args.wandb:
|
||||
cfg.use_wandb = True
|
||||
if args.no_tb:
|
||||
cfg.use_tb = False
|
||||
if args.gradcam:
|
||||
cfg.use_gradcam = True
|
||||
if args.profile:
|
||||
cfg.use_profiler = True
|
||||
if args.no_grad_norms:
|
||||
cfg.log_grad_norms = False
|
||||
|
||||
if args.output_dir is not None:
|
||||
cfg.output_dir = args.output_dir
|
||||
elif args.baseline:
|
||||
elif args.baseline and args.output_dir is None:
|
||||
cfg.output_dir = "out/gtauav/baseline"
|
||||
|
||||
train(cfg)
|
||||
|
||||
Reference in New Issue
Block a user