args = parser.parse_args() print(f"{gct()} : start time") random.seed(cfg.PROJ.SEED) torch.manual_seed(cfg.PROJ.SEED) np.random.seed(cfg.PROJ.SEED) print(f"{gct()} : model init") det = RFDet( cfg.TRAIN.score_com_strength, cfg.TRAIN.scale_com_strength, cfg.TRAIN.NMS_THRESH, cfg.TRAIN.NMS_KSIZE, cfg.TRAIN.TOPK, cfg.MODEL.GAUSSIAN_KSIZE, cfg.MODEL.GAUSSIAN_SIGMA, cfg.MODEL.KSIZE, cfg.MODEL.padding, cfg.MODEL.dilation, cfg.MODEL.scale_list, ) des = HardNetNeiMask(cfg.HARDNET.MARGIN, cfg.MODEL.COO_THRSH) model = Network(det, des, cfg.LOSS.SCORE, cfg.LOSS.PAIR, cfg.PATCH.SIZE, cfg.TRAIN.TOPK) print(f"{gct()} : to device") device = torch.device("cuda") model = model.to(device) resume = args.resume print(f"{gct()} : in {resume}")
if mgpu: print(f"{gct()} : Train with {torch.cuda.device_count()} GPUs") torch.backends.cudnn.deterministic = True random.seed(seed) torch.manual_seed(seed) np.random.seed(seed) ############################################################################### # Build the model ############################################################################### print(f"{gct()} : Build the model") det = RFDet( cfg.TRAIN.score_com_strength, cfg.TRAIN.scale_com_strength, cfg.TRAIN.NMS_THRESH, cfg.TRAIN.NMS_KSIZE, cfg.TRAIN.TOPK, cfg.MODEL.GAUSSIAN_KSIZE, cfg.MODEL.GAUSSIAN_SIGMA, ) des = HardNetNeiMask(cfg.HARDNET.MARGIN, cfg.MODEL.COO_THRSH) model = Network( det, des, cfg.LOSS.SCORE, cfg.LOSS.PAIR, cfg.PATCH.SIZE, cfg.TRAIN.TOPK ) if mgpu: model = torch.nn.DataParallel(model) model = model.to(device=device) ############################################################################### # Load train data ###############################################################################