def get_model(cfg, num_classes, device, logger):
if 'decouple' in cfg.NAME:
if cfg.TRAIN_STAGE == 1:
model = Network1(cfg, mode="train", num_classes=num_classes)
else:
model = Network(cfg,
mode="train",
num_classes=int(sum(num_classes) / 100) * 100)
else:
if isinstance(num_classes, list) and not cfg.MULTI_BRANCH:
model = Network1(cfg, mode="train", num_classes=num_classes)
elif isinstance(num_classes, list) and cfg.MULTI_BRANCH:
model = Network2(cfg, mode="train", num_classes=num_classes)
else:
model = Network(cfg, mode="train", num_classes=num_classes)
if cfg.BACKBONE.FREEZE == True:
model.freeze_backbone()
logger.info("Backbone has been freezed")
if cfg.CPU_MODE:
model = model.to(device)
else:
model = torch.nn.DataParallel(model).cuda()
return model
def get_model(cfg, num_classes, device, logger):
model = Network(cfg, mode="train", num_classes=num_classes)
if cfg.BACKBONE.FREEZE == True:
model.freeze_backbone()
logger.info("Backbone has been freezed")
if cfg.CPU_MODE:
model = model.to(device)
else:
model = torch.nn.DataParallel(model).cuda()
return model
def get_model(cfg, num_classes, device, logger):
model = Network(cfg, mode="train", num_classes=num_classes)
if cfg.BACKBONE.FREEZE == True:
model.freeze_backbone()
logger.info("Backbone has been freezed")
if not cfg.DATASET.GENERATE_CAM_BASED_DATASET and cfg.TRAIN.DISTRIBUTED:
if cfg.TRAIN.SYNCBN:
model = apex.parallel.convert_syncbn_model(model)
else:
model = model.cuda()
return model
