def create_model(config): input_info_list = create_input_infos(config) image_size = input_info_list[0].shape[-1] ssd_net = build_ssd(config.model, config.ssd_params, image_size, config.num_classes, config) compression_algo, ssd_net = create_compressed_model(ssd_net, config) ssd_net = compression_algo.model weights = config.get('weights') if weights: sd = torch.load(weights, map_location='cpu') load_state(ssd_net, sd) ssd_net.train() model, _ = prepare_model_for_execution(ssd_net, config) return compression_algo, model
def create_model(config): ssd_net = build_ssd(config.model, config.ssd_params, config.input_sample_size[-1], config.num_classes, config) ssd_net.to(config.device) compression_algo = create_compression_algorithm(ssd_net, config) ssd_net = compression_algo.model weights = config.get('weights') if weights: sd = torch.load(weights, map_location='cpu') load_state(ssd_net, sd) ssd_net.train() model, _ = prepare_model_for_execution(ssd_net, config) return compression_algo, model
def create_model(config: SampleConfig, resuming_model_sd: dict = None): input_info_list = create_input_infos(config.nncf_config) image_size = input_info_list[0].shape[-1] ssd_net = build_ssd(config.model, config.ssd_params, image_size, config.num_classes, config) weights = config.get('weights') if weights: sd = torch.load(weights, map_location='cpu') load_state(ssd_net, sd) ssd_net.to(config.device) compression_ctrl, compressed_model = create_compressed_model(ssd_net, config.nncf_config, resuming_model_sd) compressed_model, _ = prepare_model_for_execution(compressed_model, config) compressed_model.train() return compression_ctrl, compressed_model