# encoding: utf-8 """ @author: sherlock @contact: sherlockliao01@gmail.com """ import os import sys from torch import nn sys.path.append('../..') from fastreid.config import get_cfg from fastreid.engine import DefaultTrainer, default_argument_parser, default_setup from fastreid.utils.checkpoint import Checkpointer from fastreid.evaluation import ReidEvaluator from reduce_head import ReduceHead class Trainer(DefaultTrainer): @classmethod def build_evaluator(cls, cfg, num_query, output_folder=None): if output_folder is None: output_folder = os.path.join(cfg.OUTPUT_DIR, "inference") return ReidEvaluator(cfg, num_query) def setup(args): """ Create configs and perform basic setups. """ cfg = get_cfg() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() default_setup(cfg, args) return cfg def main(args): cfg = setup(args) if args.eval_only: cfg.defrost() cfg.MODEL.BACKBONE.PRETRAIN = False model = Trainer.build_model(cfg) model = nn.DataParallel(model) model = model.cuda() Checkpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=args.resume ) res = Trainer.test(cfg, model) return res trainer = Trainer(cfg) trainer.resume_or_load(resume=args.resume) return trainer.train() if __name__ == "__main__": args = default_argument_parser().parse_args() print("Command Line Args:", args) main(args)