YannAhlgrim
125 lines
3.8 KiB
Python
125 lines
3.8 KiB
Python
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"]
|
|
|
|
# -- 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 (Linear Probing vs Fine-Tuning)
|
|
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. Load Pre-trained Weights
|
|
if m_args["load_checkpoint"]:
|
|
load_path = os.path.join(args["logging"]["folder"], m_args["read_checkpoint"])
|
|
checkpoint = torch.load(load_path, map_location="cpu")
|
|
msg = model.encoder.load_state_dict(checkpoint["encoder"], strict=False)
|
|
logger.info(f"Loaded encoder from {load_path} with msg: {msg}")
|
|
|
|
# -- 5. 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, # No mask collator needed
|
|
)
|
|
|
|
# -- 6. 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"]
|
|
)
|
|
|
|
criterion = nn.CrossEntropyLoss().to(device)
|
|
model = DistributedDataParallel(model, device_ids=[torch.cuda.current_device()])
|
|
|
|
# -- 7. Training Loop
|
|
for epoch in range(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} [{itr}/{len(loader)}] Loss: {loss_meter.avg:.4f}"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|