def main():
if not os.path.isdir(cfg.CKPT):
mkdir_p(cfg.CKPT)
if args.cfg_file is not None:
shutil.copyfile(args.cfg_file, os.path.join(cfg.CKPT, args.cfg_file.split('/')[-1]))
assert_and_infer_cfg(make_immutable=False)
# Create model
model = Generalized_RCNN()
logging_rank(model, distributed=args.distributed, local_rank=args.local_rank)
# Create checkpointer
checkpointer = CheckPointer(cfg.CKPT, weights_path=cfg.TRAIN.WEIGHTS, auto_resume=cfg.TRAIN.AUTO_RESUME,
local_rank=args.local_rank)
# Load model or random-initialization
model = checkpointer.load_model(model, convert_conv1=cfg.MODEL.CONV1_RGB2BGR)
if cfg.MODEL.BATCH_NORM == 'freeze':
model = convert_bn2affine_model(model, merge=not checkpointer.resume)
elif cfg.MODEL.BATCH_NORM == 'sync':
model = convert_bn2syncbn_model(model)
model.to(args.device)
# Create optimizer
optimizer = Optimizer(model, cfg.SOLVER, local_rank=args.local_rank).build()
optimizer = checkpointer.load_optimizer(optimizer)
logging_rank('The mismatch keys: {}'.format(mismatch_params_filter(sorted(checkpointer.mismatch_keys))),
distributed=args.distributed, local_rank=args.local_rank)
# Create scheduler
scheduler = LearningRateScheduler(optimizer, cfg.SOLVER, start_iter=0, local_rank=args.local_rank)
scheduler = checkpointer.load_scheduler(scheduler)
# Create training dataset and loader
datasets = build_dataset(cfg.TRAIN.DATASETS, is_train=True, local_rank=args.local_rank)
train_loader = make_train_data_loader(datasets, is_distributed=args.distributed, start_iter=scheduler.iteration)
# Model Distributed
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.local_rank], output_device=args.local_rank,
)
else:
model = torch.nn.DataParallel(model)
# Build hooks
all_hooks = build_train_hooks(cfg, optimizer, scheduler, max_iter=cfg.SOLVER.MAX_ITER,
warmup_iter=cfg.SOLVER.WARM_UP_ITERS, ignore_warmup_time=False)
# Train
train(model, train_loader, optimizer, scheduler, checkpointer, all_hooks)
def initialize_model_from_cfg():
"""Initialize a model from the global cfg. Loads test-time weights and
set to evaluation mode.
"""
model = Generalized_RCNN(is_train=False)
# Load trained model
cfg.TEST.WEIGHTS = get_weights(cfg.CKPT, cfg.TEST.WEIGHTS)
load_weights(model, cfg.TEST.WEIGHTS)
if cfg.MODEL.BATCH_NORM == 'freeze':
model = convert_bn2affine_model(model)
model.eval()
model.to(torch.device(cfg.DEVICE))
return model