| |
| """Train FlexICM extension-layer C-TAIC for a multi-task scenario. |
| |
| Scenarios (paper Sec.IV.A): |
| s1: detection (base) -> instance (extension) |
| s2: semantic (base) -> panoptic (extension) |
| s3: detection (base) -> pose (extension) |
| |
| Two-stage training (paper Sec.III.B.3): |
| stage1: TAIC mode for extension task (SFMA + Task Connector) |
| stage2: condition mode (Prompt Generator + Condition Generator), needs base y_b_hat |
| |
| Example: |
| python scripts/train_ctaic.py -c configs/ctaic/s1_det_instance.yaml --stage 1 |
| python scripts/train_ctaic.py -c configs/ctaic/s1_det_instance.yaml --stage 2 |
| """ |
|
|
| 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 |
|
|
| 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 CTAIC, 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, |
| ) |
|
|
| SCENARIOS = { |
| "s1": {"base": "detection", "ext": "instance"}, |
| "s2": {"base": "semantic", "ext": "panoptic"}, |
| "s3": {"base": "detection", "ext": "pose"}, |
| } |
|
|
|
|
| def parse_args(argv): |
| parser = argparse.ArgumentParser("Train FlexICM C-TAIC") |
| parser.add_argument("-c", "--config", required=True) |
| parser.add_argument("--stage", type=int, choices=[1, 2], default=1) |
| 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) |
| |
| parser.set_defaults(config=given.config, stage=given.stage, **cfg) |
| for action in parser._actions: |
| if "--config" in action.option_strings: |
| action.required = False |
| break |
| args = parser.parse_args(remaining) |
| args.config = given.config |
| args.stage = given.stage |
| return args |
|
|
|
|
| def build_base_codec(args, device): |
| """Frozen base-layer TAIC used to provide y_b_hat.""" |
| base_task = SCENARIOS[args.scenario]["base"] |
| out_ch = TASK_META[base_task]["out_channels"] |
| base = TAIC(N=128, M=192, out_channels=out_ch).to(device) |
| if args.base_taic_checkpoint: |
| state, _ = load_checkpoint_dict(args.base_taic_checkpoint, map_location=device) |
| base.load_state_dict(state, strict=False) |
| elif args.base_codec: |
| state, _ = load_checkpoint_dict(args.base_codec, map_location=device) |
| base.load_base_codec(state, strict=False) |
| base.eval() |
| for p in base.parameters(): |
| p.requires_grad = False |
| return base |
|
|
|
|
| @torch.no_grad() |
| def encode_base_latent(base_model, images): |
| out = base_model(images) |
| return out["y_hat"] |
|
|
|
|
| def train_one_epoch(stage, ext_model, base_model, teacher, loader, optimizer, criterion, device, log_every=50): |
| ext_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) |
| if stage == 1: |
| out = ext_model(images, y_b_hat=None, use_condition=False) |
| else: |
| y_b = encode_base_latent(base_model, images) |
| out = ext_model(images, y_b_hat=y_b, use_condition=True) |
| 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"[stage{stage} {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()} |
|
|
|
|
| @torch.no_grad() |
| def validate(stage, ext_model, base_model, teacher, loader, criterion, device, align_divisor=256): |
| ext_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) |
| if stage == 1: |
| out = ext_model(x, use_condition=False) |
| else: |
| y_b = encode_base_latent(base_model, x) |
| out = ext_model(x, y_b_hat=y_b, use_condition=True) |
| 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()) |
| ext_model.train() |
| return {k: m.avg for k, m in meters.items()} |
|
|
|
|
| def main(argv): |
| args = parse_args(argv) |
| set_seed(getattr(args, "seed", 42)) |
| stage = args.stage |
| out_dir = exp_dir(args.root, f"{args.exp_name}_stage{stage}", args.quality_level) |
| setup_logger(os.path.join(out_dir, time.strftime("%Y%m%d_%H%M%S") + ".log")) |
| logging.info(f"Scenario {args.scenario}: {SCENARIOS[args.scenario]} stage={stage}") |
|
|
| os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_id) |
| device = "cuda" if args.cuda and torch.cuda.is_available() else "cpu" |
|
|
| ext_task = SCENARIOS[args.scenario]["ext"] |
| meta = TASK_META[ext_task] |
| out_channels = getattr(args, "out_channels", meta["out_channels"]) |
| align_mode = getattr(args, "align_mode", meta["align_mode"]) |
|
|
| train_tf = build_train_transform(args.patch_size) |
| if ext_task == "pose": |
| train_set = COCOWholeBodyImageDataset(args.dataset_path, "train2017", train_tf) |
| else: |
| train_set = COCOImageDataset(args.dataset_path, "train2017", train_tf) |
| val_set = COCOImageDataset(args.dataset_path, "val2017", 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"), |
| ) |
|
|
| base_model = build_base_codec(args, device) if stage == 2 else None |
|
|
| net = CTAIC(N=128, M=192, out_channels=out_channels).to(device) |
| if args.base_codec: |
| state, _ = load_checkpoint_dict(args.base_codec, map_location=device) |
| net.load_base_codec(state, strict=False) |
|
|
| if stage == 1: |
| net.freeze_for_stage1() |
| if args.taic_init: |
| state, _ = load_checkpoint_dict(args.taic_init, map_location=device) |
| net.load_taic_checkpoint(state) |
| else: |
| |
| init_ck = args.stage1_checkpoint or args.taic_init |
| if not init_ck: |
| raise ValueError("stage2 requires --stage1_checkpoint or taic_init in config") |
| state, _ = load_checkpoint_dict(init_ck, map_location=device) |
| net.load_taic_checkpoint(state) |
| net.freeze_for_stage2() |
|
|
| logging.info( |
| f"Trainable params: {sum(p.numel() for p in net.parameters() if p.requires_grad)/1e6:.3f}M" |
| ) |
|
|
| teacher = build_teacher( |
| ext_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.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) |
|
|
| best = float("inf") |
| for epoch in range(args.epochs): |
| logging.info(f"===== Stage {stage} Epoch {epoch}/{args.epochs} =====") |
| train_stats = train_one_epoch( |
| stage, net, base_model, teacher, train_loader, optimizer, criterion, device |
| ) |
| val_stats = validate(stage, net, base_model, 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, |
| "stage": stage, |
| "scenario": args.scenario, |
| "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:]) |
|
|
