|
|
|
|
|
|
|
|
|
|
|
import os |
|
|
import sys |
|
|
import random |
|
|
import time |
|
|
import warnings |
|
|
from loguru import logger |
|
|
|
|
|
import torch |
|
|
import torch.backends.cudnn as cudnn |
|
|
|
|
|
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." |
|
|
) |
|
|
|
|
|
|
|
|
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) |
|
|
|
