visualize_assign.py

#!/usr/bin/env python3
# Copyright (c) Megvii, Inc. and its affiliates.

import os
import sys
import random
import time
import warnings
from loguru import logger

import torch
import torch.backends.cudnn as cudnn
sys.path.append(sys.path[0].split('/tools')[0])

from yolox.exp import Exp, get_exp
from yolox.core import Trainer
from yolox.utils import configure_module, configure_omp
from yolox.tools.train import make_parser


class AssignVisualizer(Trainer):

    def __init__(self, exp: Exp, args):
        super().__init__(exp, args)
        self.batch_cnt = 0
        self.vis_dir = os.path.join(self.file_name, "vis")
        os.makedirs(self.vis_dir, exist_ok=True)

    def train_one_iter(self):
        iter_start_time = time.time()

        inps, targets = self.prefetcher.next()
        inps = inps.to(self.data_type)
        targets = targets.to(self.data_type)
        targets.requires_grad = False
        inps, targets = self.exp.preprocess(inps, targets, self.input_size)
        data_end_time = time.time()

        with torch.cuda.amp.autocast(enabled=self.amp_training):
            path_prefix = os.path.join(self.vis_dir, f"assign_vis_{self.batch_cnt}_")
            self.model.visualize(inps, targets, path_prefix)

        if self.use_model_ema:
            self.ema_model.update(self.model)

        iter_end_time = time.time()
        self.meter.update(
            iter_time=iter_end_time - iter_start_time,
            data_time=data_end_time - iter_start_time,
        )
        self.batch_cnt += 1
        if self.batch_cnt >= self.args.max_batch:
            sys.exit(0)

    def after_train(self):
        logger.info("Finish visualize assignment, exit...")


def assign_vis_parser():
    parser = make_parser()
    parser.add_argument("--max-batch", type=int, default=1, help="max batch of images to visualize")
    return parser


@logger.catch
def main(exp: Exp, args):
    if exp.seed is not None:
        random.seed(exp.seed)
        torch.manual_seed(exp.seed)
        cudnn.deterministic = True
        warnings.warn(
            "You have chosen to seed training. This will turn on the CUDNN deterministic setting, "
            "which can slow down your training considerably! You may see unexpected behavior "
            "when restarting from checkpoints."
        )

    # set environment variables for distributed training
    configure_omp()
    cudnn.benchmark = True

    visualizer = AssignVisualizer(exp, args)
    visualizer.train()


if __name__ == "__main__":
    configure_module()
    args = assign_vis_parser().parse_args()
    exp = get_exp(args.exp_file, args.name)
    exp.merge(args.opts)

    if not args.experiment_name:
        args.experiment_name = exp.exp_name

    main(exp, args)