tianma

Support optional fixed eval_size and update detection configs.

3482727 3 months ago

8.53 kB

#!/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:])

	#!/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:])