YannAhlgrim/wilds-ijepa
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fixed conflicts between cluster and github versions

Browse Source
This commit is contained in:
ahyanc25@reutlingen-university.de
2026-04-18 12:32:22 +02:00
parent 69d811ecb4 4c4d6713e0
commit f2129dc28a
6 changed files with 365 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2
View File
@@ -1,3 +1,5 @@
*.swp
*.swo
__pycache__
.venv/
experiment_logs/
configs/supervised_wilds_vith14_ep300.yaml
+29
View File
@@ -0,0 +1,29 @@
meta:
model_name: vit_huge
embed_dim: 1280
load_checkpoint: true
read_checkpoint: jepa-latest.pth.tar
use_bfloat16: true
num_classes: 182
data:
batch_size: 128
root_path: ./wilds_data
num_workers: 10
pin_mem: true
crop_size: 224
optimization:
optimizer: adamw # 'adamw' or 'sgd'
freeze_weights: true # true for linear probing, false for full fine-tuning
epochs: 50
lr: 0.001
weight_decay: 0.05
start_lr: 0.0001
final_lr: 1.0e-06
warmup: 5
ipe_scale: 1.0
logging:
folder: experiment_logs/supervised-vith14.224-bs.128-ep.300/
write_tag: linear_probe
main_distributed.py
+29 -27
View File
@@ -21,46 +21,43 @@ logger = logging.getLogger()
parser = argparse.ArgumentParser()
parser.add_argument("--folder", type=str, help="location to save submitit logs")
parser.add_argument(
'--folder', type=str,
help='location to save submitit logs')
"--batch-launch",
action="store_true",
help="whether fname points to a file to batch-lauch several config files",
)
parser.add_argument(
'--batch-launch', action='store_true',
help='whether fname points to a file to batch-lauch several config files')
"--fname",
type=str,
help="yaml file containing config file names to launch",
default="configs.yaml",
)
parser.add_argument("--partition", type=str, help="cluster partition to submit jobs on")
parser.add_argument(
'--fname', type=str,
help='yaml file containing config file names to launch',
default='configs.yaml')
"--nodes", type=int, default=1, help="num. nodes to request for job"
)
parser.add_argument(
'--partition', type=str,
help='cluster partition to submit jobs on')
parser.add_argument(
'--nodes', type=int, default=1,
help='num. nodes to request for job')
parser.add_argument(
'--tasks-per-node', type=int, default=1,
help='num. procs to per node')
parser.add_argument(
'--time', type=int, default=4300,
help='time in minutes to run job')
"--tasks-per-node", type=int, default=1, help="num. procs to per node"
)
parser.add_argument("--time", type=int, default=4300, help="time in minutes to run job")
class Trainer:
def __init__(self, fname='configs.yaml', load_model=None):
def __init__(self, fname="configs.yaml", load_model=None):
self.fname = fname
self.load_model = load_model
def __call__(self):
fname = self.fname
load_model = self.load_model
logger.info(f'called-params {fname}')
logger.info(f"called-params {fname}")
# -- load script params
params = None
with open(fname, 'r') as y_file:
with open(fname, "r") as y_file:
params = yaml.load(y_file, Loader=yaml.FullLoader)
logger.info('loaded params...')
logger.info("loaded params...")
pp = pprint.PrettyPrinter(indent=4)
pp.pprint(params)
@@ -69,7 +66,9 @@ class Trainer:
def checkpoint(self):
fb_trainer = Trainer(self.fname, True)
return submitit.helpers.DelayedSubmission(fb_trainer,)
return submitit.helpers.DelayedSubmission(
fb_trainer,
)
def launch():
@@ -83,7 +82,8 @@ def launch():
nodes=args.nodes,
ntasks_per_node=args.tasks_per_node,
cpus_per_task=10,
gpus_per_node=args.tasks_per_node)
gpus_per_node=args.tasks_per_node,
)
config_fnames = [args.fname]
@@ -91,7 +91,9 @@ def launch():
with executor.batch():
for cf in config_fnames:
fb_trainer = Trainer(cf)
job = executor.submit(fb_trainer,)
job = executor.submit(
fb_trainer,
)
trainers.append(fb_trainer)
jobs.append(job)
@@ -99,6 +101,6 @@ def launch():
print(job.job_id)
if __name__ == '__main__':
if __name__ == "__main__":
args = parser.parse_args()
launch()
main_distributed_supervised.py
+106
View File
@@ -0,0 +1,106 @@
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#
import argparse
import logging
import os
import pprint
import sys
import yaml
import submitit
from src.train_supervised import main as app_main
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logger = logging.getLogger()
parser = argparse.ArgumentParser()
parser.add_argument("--folder", type=str, help="location to save submitit logs")
parser.add_argument(
"--batch-launch",
action="store_true",
help="whether fname points to a file to batch-lauch several config files",
)
parser.add_argument(
"--fname",
type=str,
help="yaml file containing config file names to launch",
default="configs.yaml",
)
parser.add_argument("--partition", type=str, help="cluster partition to submit jobs on")
parser.add_argument(
"--nodes", type=int, default=1, help="num. nodes to request for job"
)
parser.add_argument(
"--tasks-per-node", type=int, default=1, help="num. procs to per node"
)
parser.add_argument("--time", type=int, default=4300, help="time in minutes to run job")
class Trainer:
def __init__(self, fname="configs.yaml", load_model=None):
self.fname = fname
self.load_model = load_model
def __call__(self):
fname = self.fname
load_model = self.load_model
logger.info(f"called-params {fname}")
# -- load script params
params = None
with open(fname, "r") as y_file:
params = yaml.load(y_file, Loader=yaml.FullLoader)
logger.info("loaded params...")
pp = pprint.PrettyPrinter(indent=4)
pp.pprint(params)
resume_preempt = False if load_model is None else load_model
app_main(args=params, resume_preempt=resume_preempt)
def checkpoint(self):
fb_trainer = Trainer(self.fname, True)
return submitit.helpers.DelayedSubmission(
fb_trainer,
)
def launch():
executor = submitit.SlurmExecutor(
folder=os.path.join(args.folder, "job_%j"), max_num_timeout=20
)
executor.update_parameters(
partition=args.partition,
mem_per_gpu="55G",
time=args.time,
nodes=args.nodes,
ntasks_per_node=args.tasks_per_node,
cpus_per_task=10,
gpus_per_node=args.tasks_per_node,
)
config_fnames = [args.fname]
jobs, trainers = [], []
with executor.batch():
for cf in config_fnames:
fb_trainer = Trainer(cf)
job = executor.submit(
fb_trainer,
)
trainers.append(fb_trainer)
jobs.append(job)
for job in jobs:
print(job.job_id)
if __name__ == "__main__":
args = parser.parse_args()
launch()
src/models/head.py
+22
View File
@@ -0,0 +1,22 @@
import torch
import torch.nn as nn
class ViTClassifier(nn.Module):
def __init__(self, encoder, num_classes, embed_dim):
super().__init__()
self.encoder = encoder
self.head = nn.Linear(embed_dim, num_classes)
nn.init.trunc_normal_(self.head.weight, std=0.01)
nn.init.zeros_(self.head.bias)
def forward(self, x):
# ViT -> (B, N, D)
features = self.encoder(x)
# Average Pool -> (B, D)
avg_embed = features.mean(dim=1)
logits = self.head(avg_embed)
return logits
src/train_supervised.py
+177
View File
@@ -0,0 +1,177 @@
import os
import sys
import yaml
import logging
import torch
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel
import torch.nn.functional as F
import numpy as np
from src.datasets.wilds import make_iwildcam
from src.helper import load_checkpoint, init_model, init_opt
from src.transforms import make_transforms
from src.models.head import ViTClassifier # Import our new class
from src.utils.distributed import init_distributed
from src.utils.logging import CSVLogger, AverageMeter
# --
log_timings = True
log_freq = 10
checkpoint_freq = 50
# --
_GLOBAL_SEED = 0
np.random.seed(_GLOBAL_SEED)
torch.manual_seed(_GLOBAL_SEED)
torch.backends.cudnn.benchmark = True
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logger = logging.getLogger()
def main(args):
# -- Init Distributed
world_size, rank = init_distributed()
device = torch.device(f"cuda:{torch.cuda.current_device()}")
# -- Extract Config Params
m_args = args["meta"]
o_args = args["optimization"]
d_args = args["data"]
l_args = args["logging"]
# -- Paths for saving
folder = l_args["folder"]
tag = l_args["write_tag"]
if not os.path.exists(folder):
os.makedirs(folder, exist_ok=True)
save_path = os.path.join(folder, f"{tag}" + "-ep{epoch}.pth.tar")
latest_path = os.path.join(folder, f"{tag}-latest.pth.tar")
# -- 1. Initialize Encoder
encoder, _ = init_model(
device=device,
patch_size=args.get("mask", {}).get("patch_size", 14),
crop_size=d_args["crop_size"],
model_name=m_args["model_name"],
)
# -- 2. Wrap in Classification Head
embed_dim = m_args.get("embed_dim")
model = ViTClassifier(encoder, m_args["num_classes"], embed_dim).to(device)
# -- 3. Handle Freezing
if o_args["freeze_weights"]:
logger.info("Freezing encoder weights (Linear Probing mode)")
for name, param in model.encoder.named_parameters():
param.requires_grad = False
else:
logger.info("Training full model (Fine-tuning mode)")
# -- 4. Optimizer Selection
params = [p for p in model.parameters() if p.requires_grad]
if o_args["optimizer"].lower() == "adamw":
optimizer = torch.optim.AdamW(
params, lr=o_args["lr"], weight_decay=o_args["weight_decay"]
)
else:
optimizer = torch.optim.SGD(
params, lr=o_args["lr"], momentum=0.9, weight_decay=o_args["weight_decay"]
)
# -- 5. Resume/Load Logic
start_epoch = 0
# Priority 1: Check if we are resuming from an interrupted run (latest-path)
# Priority 2: Check if we are loading a specific pre-trained checkpoint (m_args["load_checkpoint"])
checkpoint_to_load = None
resuming_interrupted = False
if os.path.exists(latest_path):
checkpoint_to_load = latest_path
resuming_interrupted = True
elif m_args["load_checkpoint"]:
checkpoint_to_load = os.path.join(folder, m_args["read_checkpoint"])
if checkpoint_to_load:
checkpoint = torch.load(checkpoint_to_load, map_location="cpu")
if resuming_interrupted:
# Load full state to resume exactly where we left off
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["opt"])
start_epoch = checkpoint["epoch"]
logger.info(
f"Resuming training from {checkpoint_to_load} at epoch {start_epoch}"
)
else:
# Loading just encoder weights for a fresh supervised run
msg = model.encoder.load_state_dict(checkpoint["encoder"], strict=False)
logger.info(
f"Loaded pre-trained encoder from {checkpoint_to_load} with msg: {msg}"
)
# -- 6. Data Setup
transform = make_transforms(crop_size=d_args["crop_size"])
_, loader, sampler = make_iwildcam(
transform=transform,
split="train",
batch_size=d_args["batch_size"],
root_path=d_args["root_path"],
rank=rank,
world_size=world_size,
collator=None,
)
criterion = nn.CrossEntropyLoss().to(device)
model = DistributedDataParallel(model, device_ids=[torch.cuda.current_device()])
# -- 7. Define Save Function
def save_checkpoint(epoch, current_loss):
save_dict = {
"model": model.module.state_dict(), # model.module because of DDP
"opt": optimizer.state_dict(),
"epoch": epoch,
"loss": current_loss,
"args": args,
}
if rank == 0:
torch.save(save_dict, latest_path)
if epoch % checkpoint_freq == 0:
torch.save(save_dict, save_path.format(epoch=epoch))
logger.info(f"Checkpoint saved at epoch {epoch}")
# -- 8. Training Loop
for epoch in range(start_epoch, o_args["epochs"]):
sampler.set_epoch(epoch)
model.train()
loss_meter = AverageMeter()
for itr, (imgs, labels, _) in enumerate(loader):
imgs, labels = imgs.to(device), labels.to(device)
with torch.cuda.amp.autocast(
enabled=m_args["use_bfloat16"], dtype=torch.bfloat16
):
outputs = model(imgs)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss_meter.update(loss.item())
if itr % 10 == 0 and rank == 0:
logger.info(
f"Epoch {epoch + 1} [{itr}/{len(loader)}] Loss: {loss_meter.avg:.4f}"
)
# Save at the end of every epoch
save_checkpoint(epoch + 1, loss_meter.avg)
if __name__ == "__main__":
main()