def make_model(args, config): # parse params with default values body_config = config["body"] fpn_config = config["fpn"] local_config = config["local"] transformer_config = config["transformer"] data_config = config["dataloader"] # BN + activation norm_act_static, norm_act_dynamic = norm_act_from_config(body_config) # Create backbone log_debug("Creating backbone model %s", body_config.get("arch")) body_fn = models.__dict__[body_config.get("arch")] body_params = body_config.getstruct("body_params") if body_config.get( "body_params") else {} body = body_fn(norm_act=norm_act_static, config=body_config, **body_params) if body_config.getboolean("pretrained"): arch = body_config.get("arch") # vgg with bn or without if body_config.get("arch").startswith("vgg"): if body_config["normalization_mode"] != 'off': arch = body_config.get("arch") + '_bn' # Download pre trained model log_debug("Downloading pre - trained model weights ") if body_config.get("source_url") == "cvut": if body_config.get("arch") not in model_urls_cvut: raise ValueError( " body arch not found in cvut witch source_url = pytorch") log_info("Downloading from m ", model_urls_cvut[arch]) state_dict = load_state_dict_from_url(model_urls_cvut[arch], progress=True) elif body_config.get("source_url") == "pytorch": if body_config.get("arch") not in model_urls: raise ValueError(" body arch not found in pytorch ") state_dict = load_state_dict_from_url(model_urls[arch], progress=True) else: raise ValueError(" try source_url = cvut or pytorch ") # Convert model to unified format and save it converted_model = body.convert(state_dict) folder = args.directory + "/image_net" if not path.exists(folder): log_debug("Create path to save pretrained backbones: %s ", folder) makedirs(folder) body_path = folder + "/" + arch + ".pth" log_debug("Saving pretrained backbones in : %s ", body_path) torch.save(converted_model, body_path) # Load converted weights to model body.load_state_dict(torch.load(body_path, map_location="cpu")) # Freeze modules in backbone for n, m in body.named_modules(): for mod_id in body_config.getstruct("num_frozen"): if ("mod%d" % mod_id) in n: freeze_params(m) # delete mods for mod_id in body_config.getstruct("num_frozen"): key = "mod" + str(mod_id) delattr(body, key) else: log_info("Initialize body to train from scratch") init_weights(body, body_config) # Feature pyramids if fpn_config.getboolean("fpn"): # Create FPN body_channels = body_config.getstruct("out_channels") fpn_inputs = fpn_config.getstruct("inputs") fpn = FPN([body_channels[inp] for inp in fpn_inputs], fpn_config.getint("out_channels"), fpn_config.getint("extra_scales"), norm_act_static, fpn_config.get("interpolation")) body = FPNBody(body, fpn, fpn_inputs) output_dim = fpn_config.getint("out_channels") else: output_dim = OUTPUT_DIM[body_config.get("arch")] # Create Local features local_gen = DenseFeatureGenerator( max_predictions=local_config.getint("max_kpt"), nms_threshold=None, score_threshold=None) local_loss = DenseFeatureLoss( name=local_config.get("loss"), gamma=local_config.getfloat("gamma"), epipolar_margin=local_config.getstruct("epipolar_margin"), cyclic_margin=local_config.getfloat("cyclic_margin")) local_pos_encoding = PositionEncodingSine(embedding_size=output_dim) local_matching_coarse = CoarseMatching( thr=local_config.getfloat("thr"), border_rm=local_config.getint("border_rm"), temperature=local_config.getfloat("temperature")) local_matching_fine = FineMatching() local_algo = DenseFeatureAlgo( generator=local_gen, loss=local_loss, coarse_matching=local_matching_coarse, fine_matching=local_matching_fine, pos_encoding=local_pos_encoding, min_level=local_config.getint("fpn_min_level"), fpn_levels=local_config.getint("fpn_levels"), win_size=local_config.getint("win_size")) local_head_coarse = DenseHead( dim=output_dim, embedding_size=transformer_config.getint("embedding_size"), num_head=transformer_config.getint("num_head"), layer_names=transformer_config.getstruct("layer_names"), layer_num=transformer_config.getint("layer_num_coarse"), attention=transformer_config.get("attention")) local_head_fine = DenseHead( dim=output_dim, embedding_size=transformer_config.getint("embedding_size"), num_head=transformer_config.getint("num_head"), layer_names=transformer_config.getstruct("layer_names"), layer_num=transformer_config.getint("layer_num_fine"), attention=transformer_config.get("attention")) if local_config.getint("embedding_size"): output_dim = local_config.getint("embedding_size") # Create a generic Local features network net = DenseNet(body, local_algo, local_head_coarse, local_head_fine) return net, output_dim
def make_model(args, config): # parse params with default values body_config = config["body"] fpn_config = config["fpn"] ir_config = config["global"] data_config = config["dataloader"] # get output dimensionality size net_modules = [] # BN + activation norm_act_static, norm_act_dynamic = norm_act_from_config(body_config) # Create backbone log_debug("Creating backbone model %s", body_config.get("arch")) body_fn = models.__dict__[body_config.get("arch")] body_params = body_config.getstruct("body_params") if body_config.get( "body_params") else {} body = body_fn(norm_act=norm_act_static, config=body_config, **body_params) net_modules.append("body") if body_config.getboolean("pretrained"): arch = body_config.get("arch") # vgg with bn or without if body_config.get("arch").startswith("vgg"): if body_config["normalization_mode"] != 'off': arch = body_config.get("arch") + '_bn' # Download pre trained model log_debug("Downloading pre - trained model weights ") if body_config.get("source_url") == "cvut": if body_config.get("arch") not in model_urls_cvut: raise ValueError( " body arch not found in cvut witch source_url = pytorch") log_info("Downloading from m ", model_urls_cvut[arch]) state_dict = load_state_dict_from_url(model_urls_cvut[arch], progress=True) elif body_config.get("source_url") == "pytorch": if body_config.get("arch") not in model_urls: raise ValueError(" body arch not found in pytorch ") state_dict = load_state_dict_from_url(model_urls[arch], progress=True) else: raise ValueError(" try source_url = cvut or pytorch ") # Convert model to unified format and save it converted_model = body.convert(state_dict) folder = args.directory + "/image_net" if not path.exists(folder): log_debug("Create path to save pretrained backbones: %s ", folder) makedirs(folder) body_path = folder + "/" + arch + ".pth" log_debug("Saving pretrained backbones in : %s ", body_path) torch.save(converted_model, body_path) # Load converted weights to model body.load_state_dict(torch.load(body_path, map_location="cpu")) # Freeze modules in backbone for n, m in body.named_modules(): for mod_id in range(1, body_config.getint("num_frozen") + 1): if ("mod%d" % mod_id) in n: freeze_params(m) else: log_info("Initialize body to train from scratch") init_weights(body, body_config) # Feature pyramids if fpn_config.getboolean("fpn"): # Create FPN body_channels = body_config.getstruct("out_channels") fpn_inputs = fpn_config.getstruct("inputs") fpn_outputs = fpn_config.getstruct("outputs") fpn = FPN([body_channels[inp] for inp in fpn_inputs], fpn_config.getint("out_channels"), fpn_config.getint("extra_scales"), norm_act_static, fpn_config.get("interpolation")) body = FPNBody(body, fpn, fpn_inputs) output_dim = fpn_config.getint("out_channels") else: output_dim = OUTPUT_DIM[body_config.get("arch")] # Create Image Retrieval global_loss = globalFeatureLoss(name=ir_config.get("loss"), sigma=ir_config.getfloat("loss_margin")) global_algo = globalFeatureAlgo( loss=global_loss, min_level=ir_config.getint("fpn_min_level"), fpn_levels=ir_config.getint("fpn_levels")) global_head = globalHead(pooling=ir_config.getstruct("pooling"), normal=ir_config.getstruct("normal"), dim=output_dim) # Data augmentation augment = RandomAugmentation(rgb_mean=data_config.getstruct("rgb_mean"), rgb_std=data_config.getstruct("rgb_std")) # Create a generic image retrieval network net = ImageRetrievalNet(body, global_algo, global_head, augment=augment) return net, net_modules, output_dim