(args)
| 20 | |
| 21 | |
| 22 | def main(args): |
| 23 | cfg = get_default_cfg() |
| 24 | if args.cfg_file: |
| 25 | cfg.merge_from_file(args.cfg_file) |
| 26 | cfg.merge_from_list(args.opts) |
| 27 | cfg.freeze() |
| 28 | |
| 29 | device = torch.device(cfg.DEVICE) |
| 30 | if cfg.SEED >= 0: |
| 31 | set_random_seed(cfg.SEED) |
| 32 | |
| 33 | print("Creating model...") |
| 34 | model = COAT(cfg) |
| 35 | model.to(device) |
| 36 | |
| 37 | print("Loading data...") |
| 38 | train_loader = build_train_loader(cfg) |
| 39 | gallery_loader, query_loader = build_test_loader(cfg) |
| 40 | |
| 41 | softmax_criterion_s2 = None |
| 42 | softmax_criterion_s3 = None |
| 43 | if cfg.MODEL.LOSS.USE_SOFTMAX: |
| 44 | softmax_criterion_s2 = SoftmaxLoss(cfg) |
| 45 | softmax_criterion_s3 = SoftmaxLoss(cfg) |
| 46 | softmax_criterion_s2.to(device) |
| 47 | softmax_criterion_s3.to(device) |
| 48 | |
| 49 | if args.eval: |
| 50 | assert args.ckpt, "--ckpt must be specified when --eval enabled" |
| 51 | resume_from_ckpt(args.ckpt, model) |
| 52 | evaluate_performance( |
| 53 | model, |
| 54 | gallery_loader, |
| 55 | query_loader, |
| 56 | device, |
| 57 | use_gt=cfg.EVAL_USE_GT, |
| 58 | use_cache=cfg.EVAL_USE_CACHE, |
| 59 | use_cbgm=cfg.EVAL_USE_CBGM, |
| 60 | gallery_size=cfg.EVAL_GALLERY_SIZE, |
| 61 | ) |
| 62 | exit(0) |
| 63 | |
| 64 | params = [p for p in model.parameters() if p.requires_grad] |
| 65 | if cfg.MODEL.LOSS.USE_SOFTMAX: |
| 66 | params_softmax_s2 = [p for p in softmax_criterion_s2.parameters() if p.requires_grad] |
| 67 | params_softmax_s3 = [p for p in softmax_criterion_s3.parameters() if p.requires_grad] |
| 68 | params.extend(params_softmax_s2) |
| 69 | params.extend(params_softmax_s3) |
| 70 | |
| 71 | optimizer = torch.optim.SGD( |
| 72 | params, |
| 73 | lr=cfg.SOLVER.BASE_LR, |
| 74 | momentum=cfg.SOLVER.SGD_MOMENTUM, |
| 75 | weight_decay=cfg.SOLVER.WEIGHT_DECAY, |
| 76 | ) |
| 77 | |
| 78 | lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( |
| 79 | optimizer, milestones=cfg.SOLVER.LR_DECAY_MILESTONES, gamma=cfg.SOLVER.GAMMA |
no test coverage detected