#!/usr/bin/env python3 """Train FlexICM base-layer TAIC for one of the five machine tasks. Example: python scripts/train_taic.py -c configs/taic/detection.yaml """ from __future__ import annotations import argparse import logging import os import sys import time from datetime import datetime import torch from torch.utils.data import DataLoader # allow running from repo root REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) if REPO_ROOT not in sys.path: sys.path.insert(0, REPO_ROOT) from flexicm.data import ( COCOImageDataset, COCOWholeBodyImageDataset, build_train_transform, build_test_transform, ) from flexicm.models import TAIC from flexicm.tasks import TASK_META, build_teacher from flexicm.tasks.losses import TAICCriterion from flexicm.utils.alignment import Alignment from flexicm.utils.train_utils import ( AverageMeter, CustomDataParallel, adamw_trainable, exp_dir, load_checkpoint_dict, load_yaml_config, save_checkpoint, set_seed, setup_logger, ) def parse_args(argv): parser = argparse.ArgumentParser("Train FlexICM TAIC") parser.add_argument("-c", "--config", required=True, help="YAML config path") parser.add_argument("--name", default=datetime.now().strftime("%Y-%m-%d_%H_%M_%S")) given, remaining = parser.parse_known_args(argv) cfg = load_yaml_config(given.config) # -c already consumed by the first parse; keep it as a default for the second pass parser.set_defaults(config=given.config, **cfg) for action in parser._actions: if "--config" in action.option_strings: action.required = False break parser.add_argument("-T", "--TEST", action="store_true") args = parser.parse_args(remaining) args.config = given.config return args @torch.no_grad() def validate(model, teacher, loader, criterion, device, align_divisor=256): model.eval() meters = {k: AverageMeter() for k in ("loss", "bpp", "distortion")} for images in loader: images = images.to(device) align = Alignment(divisor=align_divisor, mode="pad", padding_mode="constant").to(device) x = align.align(images) out = model(x) # Keep teacher + codec features on the aligned (padded) grid so the Swin # backbone always sees a patch/window-divisible size. bpp still uses the # original pixel count below. with torch.no_grad(): gt = teacher.gt_features(x) pred = teacher.pred_features(out["h"]) N, _, H, W = images.shape stats = criterion(out, pred, gt, num_pixels=N * H * W) for k in meters: meters[k].update(stats[k].item()) model.train() return {k: m.avg for k, m in meters.items()} def train_one_epoch(model, teacher, loader, optimizer, criterion, device, log_every=50): model.train() teacher.eval() meters = {k: AverageMeter() for k in ("loss", "bpp", "distortion")} for i, images in enumerate(loader): images = images.to(device) optimizer.zero_grad(set_to_none=True) out = model(images) with torch.no_grad(): gt = teacher.gt_features(images) pred = teacher.pred_features(out["h"]) N, _, H, W = images.shape stats = criterion(out, pred, gt, num_pixels=N * H * W) stats["loss"].backward() optimizer.step() for k in meters: meters[k].update(stats[k].item(), n=images.size(0)) if i % log_every == 0: logging.info( f"[{i}/{len(loader)}] loss={meters['loss'].avg:.4f} " f"bpp={meters['bpp'].avg:.4f} D={meters['distortion'].avg:.6f}" ) return {k: m.avg for k, m in meters.items()} def main(argv): args = parse_args(argv) set_seed(getattr(args, "seed", 42)) out_dir = exp_dir(args.root, args.exp_name, args.quality_level) setup_logger(os.path.join(out_dir, time.strftime("%Y%m%d_%H%M%S") + ".log")) logging.info(f"Config: {args.config}") for k, v in sorted(vars(args).items()): logging.info(f"{k}: {v}") os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_id) device = "cuda" if args.cuda and torch.cuda.is_available() else "cpu" task = args.task meta = TASK_META[task] out_channels = getattr(args, "out_channels", meta["out_channels"]) align_mode = getattr(args, "align_mode", meta["align_mode"]) # data train_tf = build_train_transform(args.patch_size) if task == "pose": train_set = COCOWholeBodyImageDataset(args.dataset_path, "val2017", train_tf) else: train_set = COCOImageDataset(args.dataset_path, "val2017", train_tf) val_split = getattr(args, "val_split", "val2017") val_root = os.path.join(args.dataset_path, val_split) if not os.path.isdir(val_root): val_root = os.path.join(args.dataset_path, "Kodak") if os.path.isdir( os.path.join(args.dataset_path, "Kodak") ) else os.path.join(args.dataset_path, "train2017") val_set = COCOImageDataset( os.path.dirname(val_root), os.path.basename(val_root), build_test_transform() ) if os.path.basename(val_root) in ("train2017", "val2017") else __import__( "flexicm.data", fromlist=["ImageFolderDataset"] ).ImageFolderDataset(val_root, build_test_transform()) train_loader = DataLoader( train_set, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=(device == "cuda"), drop_last=True, ) val_loader = DataLoader( val_set, batch_size=args.test_batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=(device == "cuda"), ) # models net = TAIC(N=128, M=192, out_channels=out_channels).to(device) if args.base_codec: logging.info(f"Loading base TIC codec from {args.base_codec}") state, _ = load_checkpoint_dict(args.base_codec, map_location=device) net.load_base_codec(state, strict=False) if getattr(args, "freeze_base_codec", True): net.freeze_base_codec() logging.info("Base TIC codec frozen") else: for p in net.parameters(): p.requires_grad = True logging.info("Base TIC codec unfrozen; training the entire TAIC model") logging.info( f"Trainable params: {sum(p.numel() for p in net.parameters() if p.requires_grad)/1e6:.3f}M / " f"total {sum(p.numel() for p in net.parameters())/1e6:.3f}M" ) teacher = build_teacher( task, pretrained_backbone=getattr(args, "pretrained_backbone", True), use_official_teacher=getattr(args, "use_official_teacher", True), task_config=getattr(args, "task_config", None), task_checkpoint=getattr(args, "task_checkpoint", None), device=device, ) teacher = teacher.to(device).eval() if args.checkpoint: logging.info(f"Resume/load {args.checkpoint}") state, raw = load_checkpoint_dict(args.checkpoint, map_location=device) net.load_state_dict(state, strict=False) if args.cuda and torch.cuda.device_count() > 1: net = CustomDataParallel(net) optimizer = adamw_trainable(net, lr=args.learning_rate) criterion = TAICCriterion(lmbda=args.lmbda, align_mode=align_mode) if args.TEST: stats = validate(net, teacher, val_loader, criterion, device) logging.info(f"TEST {stats}") return best = float("inf") for epoch in range(args.epochs): logging.info(f"===== Epoch {epoch}/{args.epochs} =====") train_stats = train_one_epoch(net, teacher, train_loader, optimizer, criterion, device) if epoch % 2 == 0: val_stats = validate(net, teacher, val_loader, criterion, device) logging.info(f"train={train_stats} val={val_stats}") is_best = val_stats["loss"] < best best = min(best, val_stats["loss"]) if args.save: save_checkpoint( { "epoch": epoch, "state_dict": net.module.state_dict() if hasattr(net, "module") else net.state_dict(), "optimizer": optimizer.state_dict(), "loss": val_stats["loss"], "args": vars(args), }, is_best, out_dir, ) if __name__ == "__main__": main(sys.argv[1:])