def selective_transfer(self, image_batch, mask_batch, targets, step): model_scores = dict(zip(self.available_surrogates_list, [0] * len(self.available_surrogates_list))) mse_criterion = torch.nn.MSELoss(reduction='mean') batch_indices = torch.arange(image_batch.size(0)) step.eps = self.args_dict['sigma'] x = image_batch.clone().detach().requires_grad_(False) if self.args_dict['targeted']: original_labels = torch.argmax(predict(self.model, x), dim=1) else: original_labels = targets x = torch.cat([x.unsqueeze(0)] * self.args_dict['num_transformations']) x = step.random_perturb(x, mask_batch) predictions = [] labels = [] for current_x in x: predictions.append(predict(self.model, current_x)) current_labels = torch.argmax(predictions[-1], dim=1) labels.append(current_labels) self.args_dict['label_shifts'] += torch.sum(~torch.eq(current_labels, original_labels)).item() for arch in self.available_surrogates_list: current_model = get_model(arch, 'standard', freeze=True, device=self.args_dict['device']).eval() for index, current_x in enumerate(x): current_predictions = predict(current_model, current_x) current_loss = mse_criterion(current_predictions[batch_indices, labels[index]], predictions[index][batch_indices, labels[index]]) model_scores[arch] += current_loss.item() to_device(current_model, 'cpu') surrogates_list = [arch for arch in sorted(model_scores, key=model_scores.get) [:self.args_dict['num_surrogates']]] if self.args_dict['similarity_coeffs']: scores_reversed = torch.FloatTensor([model_scores[arch] for arch in surrogates_list][::-1]) coeffs = (scores_reversed / torch.sum(scores_reversed)).tolist() else: coeffs = [1 / len(surrogates_list)] * len(surrogates_list) self.similarity_coeffs = (dict(zip(surrogates_list, coeffs))) ALL_SIMILARITY_COEFFS.append(self.similarity_coeffs) surrogate_models = [get_model(arch, pretrained=True, freeze=True).eval() for arch in surrogates_list] return surrogate_models
def predict( workdir: pathlib.Path, config: ml_collections.ConfigDict, output_filepath: str, ) -> None: """Generates model predictions using the best available checkpoint.""" # Set seed for reproducibility. tf.random.set_seed(config.seed) if config.train.get('use_mixed_precision', False): tf.keras.mixed_precision.set_global_policy('mixed_float16') # Load the best checkpoint config, model = model_utils.get_model(config) best_checkpoint = tf.train.latest_checkpoint(workdir / 'checkpoints') print(f'Loading model from {best_checkpoint}...') model.load_weights(best_checkpoint).expect_partial() predict_ds = input_pipeline.build_predict_dataset(config.dataset, config.outcomes, cache=config.dataset.get( 'use_cache', False)) batched_predictions = predict_utils.generate_predictions(model, predict_ds) predictions_df = predict_utils.merge_batched_predictions( batched_predictions) # Write the DataFrame to a CSV. with open(output_filepath, mode='wt') as f: predictions_df.to_csv(f, sep=',', index=False)
def _restore(self, checkpoint_path): if hasattr(self, 'device'): checkpoint = torch.load(checkpoint_path, self.device) else: checkpoint = torch.load(checkpoint_path) # saved_model = torch.load(checkpoint_path + '.args') # self.model = model_utils.get_model(saved_model['args']) self.model = model_utils.get_model(checkpoint['model']['args']) self.model.to(self.device) self.model.load_state_dict(checkpoint['model']['state']) permutation_params = filter( lambda p: hasattr(p, '_is_perm_param') and p._is_perm_param, self.model.parameters()) unstructured_params = filter( lambda p: not (hasattr(p, '_is_perm_param') and p._is_perm_param), self.model.parameters()) self.optimizer = optim.Adam([{ 'params': permutation_params }, { 'params': unstructured_params }], ) # self.optimizer = optim.Adam([{'params': permutation_params, 'weight_decay': 0.0, 'lr': config['plr']}, # {'params': unstructured_params}], # lr=config['lr'], weight_decay=config['weight_decay']) self.optimizer.load_state_dict(checkpoint['optimizer']) # self.optimizer.param_groups[1].update({'weight_decay': 0.0}) # self.optimizer.param_groups[0].update({'params': permutation_params}) # self.optimizer.param_groups[1].update({'params': unstructured_params}) # self.scheduler = optim.lr_scheduler.StepLR(self.optimizer) self.scheduler.load_state_dict(checkpoint['scheduler']) self.scheduler.optimizer = self.optimizer
def main(args): target_model_name = args["model"] model, _, _ = get_model( target_model_name) model.summary()
def __init__(self, training_args_dict, pgd_args_dict, criterion=torch.nn.CrossEntropyLoss(), optimizer=torch.optim.Adam): if training_args_dict['checkpoint_location'] is not None: self.model = load_model( location=training_args_dict['checkpoint_location']) training_args_dict['arch'] = self.model.arch else: self.model = get_model( arch=training_args_dict['arch'], pretrained=(True if training_args_dict['pretrained'] else None)) self.normalize = Normalizer(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) self.training_args_dict = training_args_dict self.pgd_args_dict = pgd_args_dict self.adversarial = training_args_dict['adversarial'] self.attacker = None self.criterion = criterion self.optimizer = optimizer(self.model.parameters(), lr=training_args_dict['learning_rate']) self.losses = []
def _setup(self, config): self.config = config device = config['device'] self.device = device torch.manual_seed(config['seed']) if self.device == 'cuda': torch.cuda.manual_seed(config['seed']) # model self.model = model_utils.get_model(config['model']).to(device) self.model_args = config['model'] # count parameters self.nparameters = sum(param.nelement() for param in self.model.parameters()) print("Parameter count: ", self.nparameters) # dataset self.train_loader, self.valid_loader, self.test_loader = dataset_utils.get_dataset(config['dataset']) structured_params = filter(lambda p: hasattr(p, '_is_structured') and p._is_structured, self.model.parameters()) unstructured_params = filter(lambda p: not (hasattr(p, '_is_structured') and p._is_structured), self.model.parameters()) if config['optimizer'] == 'Adam': self.optimizer = optim.Adam([{'params': structured_params, 'weight_decay': 0.0}, {'params': unstructured_params}], lr=config['lr'], weight_decay=config['weight_decay']) else: self.optimizer = optim.SGD([{'params': structured_params, 'weight_decay': 0.0}, {'params': unstructured_params}], lr=config['lr'], momentum=0.9, weight_decay=config['weight_decay']) # scheduler if config['lr_decay']['milestones'] is not None: self.scheduler = optim.lr_scheduler.MultiStepLR(self.optimizer, milestones=config['lr_decay']['milestones'], gamma=config['lr_decay']['factor']) else: self.scheduler = optim.lr_scheduler.StepLR(self.optimizer, step_size=config['lr_decay']['period'], gamma=config['lr_decay']['factor']) self.switch_ams = config['switch_ams']
def evaluate( workdir: pathlib.Path, config: ml_collections.ConfigDict, split: input_pipeline.Split, ) -> None: """Evaluates the model using the best available checkpoint.""" # Set seed for reproducibility. tf.random.set_seed(config.seed) if config.train.get('use_mixed_precision', False): tf.keras.mixed_precision.set_global_policy('mixed_float16') config, model = model_utils.get_model(config) dataset = input_pipeline.build_dataset(split, config.dataset, config.outcomes, cache=config.dataset.get( 'use_cache', False)) # Load the best checkpoint and run validation. best_checkpoint = tf.train.latest_checkpoint(workdir / 'checkpoints') print(f'Loading model from {best_checkpoint}...') model.load_weights(best_checkpoint).expect_partial() model.evaluate(dataset, verbose=1, return_dict=True)
def __init__(self, model, args_dict, attack_step=LinfStep, mask_batch=None): self.model = model self.args_dict = args_dict self.similarity_coeffs = {} if args_dict['transfer']: self.loss = self.transfer_loss self.available_surrogates_list = copy.copy(ARCHS_LIST) self.available_surrogates_list.remove(args_dict['arch']) if not args_dict['selective']: surrogates_list = random.sample(self.available_surrogates_list, args_dict['num_surrogates']) coeffs = [1 / len(surrogates_list)] * len(surrogates_list) self.similarity_coeffs = (dict(zip(surrogates_list, coeffs))) ALL_SIMILARITY_COEFFS.append(self.similarity_coeffs) self.surrogate_models = [get_model(arch, pretrained=True, freeze=True).eval() for arch in surrogates_list] else: self.args_dict['label_shifts'] = 0 else: self.loss = self.normal_loss if args_dict['norm'] == 'l2': attack_step = L2Step self.criterion = torch.nn.CrossEntropyLoss() self.optimization_direction = -1 if args_dict['unadversarial'] or args_dict['targeted'] else 1 self.mask_batch = mask_batch self.attack_step = attack_step
def main(params): use_embedding_matrix = not params['not_use_embedding_matrix'] ans_to_id, id_to_ans = get_most_common_answers(vqa_train, vqa_val, int(params['num_answers']), params['ans_types'], show_top_ans=False, use_test=params['use_test']) embedding_matrix, ques_train_map, ans_train_map, img_train_map, ques_to_img_train = \ get_train_data(params['ans_types'], params['use_test'], use_embedding_matrix) _, ques_val_map, ans_val_map, img_val_map, ques_to_img_val = get_val_data(params['ans_types'], params['use_test'], use_embedding_matrix) normalize_image_embeddings([img_train_map, img_val_map]) if not params['use_test']: filtered_ann_ids_train = set(vqa_train.getQuesIds(ansTypes=params['ans_types'])) filtered_ann_ids_val = set(vqa_val.getQuesIds(ansTypes=params['ans_types'])) else: filtered_ann_ids_train = set(vqa_train.getQuesIds(ansTypes=params['ans_types']) + vqa_val.getQuesIds(ansTypes=params['ans_types'])) ques_train_ids = ques_train_map.keys() ques_val_ids = ques_val_map.keys() ques_train_ids = np.array([i for i in ques_train_ids if i in filtered_ann_ids_train]) if not params['use_test']: ques_val_ids = np.array([i for i in ques_val_ids if i in filtered_ann_ids_val]) train_dim, val_dim = len(ques_train_ids), len(ques_val_ids) print "Loaded dataset with train size %d and val size %d" % (train_dim, val_dim) if not params['eval_only']: train_model(ques_train_map, ans_train_map, img_train_map, ques_train_ids, ques_to_img_train, ques_val_map, ans_val_map, img_val_map, ques_val_ids, ques_to_img_val, id_to_ans, train_dim, val_dim, params['ans_types'], params, embedding_matrix) else: savedir = "models/%s_%s" % (params['model'], str(params['num_answers'])) print "Loading model" model = get_model( dropout_rate=float(params['dropout_rate']), regularization_rate=float(params['regularization_rate']), embedding_size=int(params['embedding_size']), num_classes=int(params['num_answers']), model_name=params['model'], embedding_matrix=embedding_matrix) model.load_weights("%s/%s_epoch_%d_weights.h5" % (savedir, params['model'], params['eval_epoch'])) if params['use_test']: evaluate_for_test( quesFile_test, model, params['batch_size'], ques_val_map, img_val_map, id_to_ans, params['embedding_size'], params['use_first_words'], use_embedding_matrix) else: evaluate(model, vqa_val, params['batch_size'], ques_val_map, img_val_map, id_to_ans, params['embedding_size'], params['use_first_words'], use_embedding_matrix, verbose=True)
def load_model( checkpoint_dir: pathlib.Path, config: ml_collections.ConfigDict ) -> Tuple[ml_collections.ConfigDict, tf.keras.Model]: """Loads a model using the best checkpoint at the given path.""" _, model = model_utils.get_model(config) best_checkpoint = tf.train.latest_checkpoint(checkpoint_dir) print(f'Loading model from {best_checkpoint}...') model.load_weights(best_checkpoint).expect_partial() return model
def main(): parser = argparse.ArgumentParser() parser.add_argument('--arch', type=str, choices=ARCHS_LIST, default='resnet50') parser.add_argument('--dataset', type=str, default='dataset/imagenet-airplanes.pt') parser.add_argument('--pretrained', default=False, action='store_true') parser.add_argument('--checkpoint_location', type=str, default=None) parser.add_argument('--epochs', type=int, default=1) parser.add_argument('--learning_rate', type=float, default=1e-2) parser.add_argument('--weight_averaging', default=False, action='store_true') parser.add_argument('--adversarial', default=False, action='store_true') parser.add_argument('--save_file_location', type=str, default='models/' + str(get_current_time()) + '.pt') args_dict = vars(parser.parse_args()) validate_save_file_location(args_dict['save_file_location']) if os.path.exists(args_dict['dataset']): dataset_properties = torch.load(args_dict['dataset']) pgd_args_dict = PGD_DEFAULT_ARGS_DICT pgd_args_dict['arch'] = args_dict['arch'] pgd_args_dict['dataset'] = dataset_properties['images'] pgd_args_dict['eps'] = 32 / 255.0 pgd_args_dict['step_size'] = 32 / 255.0 images = torch.load(dataset_properties['images']) if dataset_properties['labels'] is None: eval_model = get_model(arch=args_dict['arch'], pretrained=True) normalize = Normalizer(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) labels = [ torch.argmax(eval_model(normalize(x.unsqueeze(0)))) for x in images ] else: labels = torch.load(dataset_properties['labels']) trainer = Trainer(args_dict, pgd_args_dict) trainer.fit(images, labels) trainer.serialize() else: raise ValueError('Specified dataset location is incorrect!')
def get_component_network(template_file=None, binary_linear_file=None, freeze_kernels=False): model = mutil.get_model(args.arch) if template_file is not None: model.load_state_dict(torch.load(template_file, map_location='cpu')) if freeze_kernels: mutil.freeze_model_parameters_(model) mutil.set_module_trainable_(model, torch.nn.BatchNorm2d) model.linear = torch.nn.Linear(model.linear.in_features, 2) if binary_linear_file is not None: model.linear.load_state_dict( torch.load(binary_linear_file, map_location='cpu')) return model
def _restore(self, checkpoint_path): if hasattr(self, 'device'): checkpoint = torch.load(checkpoint_path, self.device) else: checkpoint = torch.load(checkpoint_path) self.model = model_utils.get_model(checkpoint['model']['args']) self.model.to(self.device) self.model.load_state_dict(checkpoint['model']['state']) # TODO: refactor this into an optimizer constructing helper structured_params = filter(lambda p: hasattr(p, '_is_structured') and p._is_structured, self.model.parameters()) unstructured_params = filter(lambda p: not (hasattr(p, '_is_structured') and p._is_structured), self.model.parameters()) self.optimizer = optim.Adam([{'params': structured_params}, {'params': unstructured_params}],) self.optimizer.load_state_dict(checkpoint['optimizer']) self.scheduler.load_state_dict(checkpoint['scheduler']) self.scheduler.optimizer = self.optimizer
def create_adversarial_dataset(results_location): results = torch.load(results_location) dataset = torch.load(results['args_dict']['dataset']) folder_location = 'dataset/adversarial/' + results['args_dict']['save_file_location'].split('/')[-1][:-3] if hasattr(dataset, 'category'): folder_location = os.path.join(folder_location, dataset.category) images_location = os.path.join(folder_location, 'images') masks_location = os.path.join(folder_location, 'masks') if not os.path.exists(folder_location): os.makedirs(folder_location) os.makedirs(images_location) if results['args_dict']['masks']: os.makedirs(masks_location) for i in range(0, len(results['adversarial_examples'])): adversarial_example = results['adversarial_examples'][i] save_image(adversarial_example, os.path.join(images_location, str(i) + '.png')) if results['args_dict']['masks']: _, mask = dataset[i] save_image(mask, os.path.join(masks_location, str(i) + '.png')) if results['args_dict']['masks']: transform = torchvision.transforms.ToTensor() parent_directory = os.path.abspath(folder_location + '/../') adversarial_dataset = datasets.CocoCategory(location=parent_directory, category=dataset.category, transform=transform) torch.save(adversarial_dataset, os.path.join(parent_directory, 'images.pt')) with open(os.path.join(parent_directory, 'args_dict.json'), 'w') as file: json.dump(results['args_dict'], file) else: model = get_model('resnet50', 'standard').eval() preprocessor = ImageNetPreprocessor(location=images_location, model=model, resize=False) preprocessor.set_dataset_images() adversarial_dataset = preprocessor.dataset_images torch.save(adversarial_dataset, os.path.join(folder_location, 'images.pt')) with open(os.path.join(folder_location, 'args_dict.json'), 'w') as file: json.dump(results['args_dict'], file)
def main(args): target_model_name = args["model"] target_layer_name = args["layer"] input_img_path = args["img_path"] model, preprocess_input, decode_predictions = get_model(target_model_name) model.summary() heatmap = make_heatmap(model, preprocess_input, decode_predictions, target_layer_name, input_img_path) superimposed_img, heatmap, img = superimpose(input_img_path, heatmap) #---------------------- # plot imgs #---------------------- plt.figure(figsize=(10, 5)) # figsize=(w,h) plt.subplot(1, 3, 1) # input img plt.title('img') plt.imshow(img) plt.axis('off') # heatmap plt.subplot(1, 3, 2) plt.title('heatmap') plt.imshow(heatmap, cmap='hot') plt.axis('off') # plt.colorbar() # heatmap + img plt.subplot(1, 3, 3) plt.title('heatmap+img') plt.imshow(np.uint8(superimposed_img)) plt.axis('off') basename = os.path.basename(input_img_path) basename, _ = os.path.splitext(basename) plt.savefig(os.path.join(os.getcwd(), basename + "_gradcam.png")) plt.show()
def main(): time = get_current_time() parser = argparse.ArgumentParser() parser.add_argument('--arch', type=str, choices=ARCHS_LIST, default='resnet50') parser.add_argument('--pretrained', default=False, action='store_true') parser.add_argument('--checkpoint_location', type=str, default=None) parser.add_argument('--from_robustness', default=False, action='store_true') parser.add_argument('--dataset', type=str, default='dataset/coco') parser.add_argument('--normalize_grads', default=False, action='store_true') parser.add_argument('--save_file_location', type=str, default='results/gradient/' + time + '.pt') args_dict = vars(parser.parse_args()) validate_save_file_location(args_dict['save_file_location']) if args_dict['checkpoint_location'] is not None: model = load_model( location=args_dict['checkpoint_location'], arch=args_dict['arch'], from_robustness=args_dict['from_robustness']).cuda().eval() else: model = get_model( args_dict['arch'], True if [args_dict['pretrained']] else False).cuda().eval() criterion = torch.nn.CrossEntropyLoss(reduction='none') averages = get_averages_dict(model, criterion, args_dict) torch.save({ 'averages': averages, 'args': args_dict }, args_dict['save_file_location'])
def main(): MODELS_LIST_WITH_ALL = MODELS_LIST.append('all') parser = argparse.ArgumentParser() parser.add_argument('--csv_location', type=str, required=True) parser.add_argument('--model', type=str, choices=MODELS_LIST_WITH_ALL, default='all') args = parser.parse_args() if os.path.exists( args.csv_location) and args.csv_location.endswith('.csv'): if args.model != 'all': model = get_model(args.model) score = compare_predictions(args.csv_location, model) print('Score for model ' + args.model + ' is ' + str(score)) else: summarize_predictions_scores(args.csv_location) else: raise ValueError('Selected path is not a folder')
def main(): parser = argparse.ArgumentParser() parser.add_argument('-c', '--config', default='config/config.yml', help='path to config file') parser.add_argument('-d', '--debug', action='store_true', help='run in debug mode') args = parser.parse_args() global debug debug = args.debug with open(args.config) as f: config = yaml.load(f) print(yaml.dump(config)) random.seed(config['seed']) appenders = {} data_path = config['log_path'] appenders[POWER_REFS_LOG] = os.path.join(data_path, 'prefs_log.csv') appenders[STATES_LOG] = os.path.join(data_path, 'states_log.csv') appenders[COMMON_LOG] = os.path.join(data_path, 'common_log') appenders[BATTERY_LOG] = os.path.join(data_path, 'battery_log.csv') appenders[OTHERS_LOG] = os.path.join(data_path, 'others_log.csv') logger = Logger(appenders) model = model_utils.get_model(TIME_QUANT, TIME_SCALE, logger, config['state_file'], config['weather_data'], 100000 if debug else None) try: model.load_state(config['state_file']) except FileNotFoundError: stderr.write('State file not found. Starting from scratch\n') run_event = threading.Event() run_event.set() elastic = Elasticsearch([ f"http://{config['elasticsearch']['auth']}@{config['elasticsearch']['host']}" ]) charge_controller = battery_controller.create_controller( config['charge_host'], config['charge_port'], data_path, "charge") discharge_controller = battery_controller.create_controller( config['discharge_host'], config['discharge_port'], data_path, "discharge") heater_controllers = [] if config['use_real_heaters']: heater_controllers = [ HeaterController(config['heaters_port'], addr, log_folder=config['log_path'], log_name="heater" + str(addr)) for addr in config['heaters_addr'] ] send = ThreadSend(run_event, elastic, model, logger, charge_controller, discharge_controller, heater_controllers) listen = ThreadListen(run_event, model, logger) if config['use_real_bat']: charge_controller.start() discharge_controller.start() send.start() listen.start() send.join() listen.join()
def train_template_network(loss='default'): """Obtain CIFAR10-trained template network. Training parameters follow original ResNet paper. Args: loss: Choose from 'default'/'sgm'/'l2' """ # Use training parameters of original ResNet paper split_index = 45000 batch_size = 128 lr = 1e-1 momentum = 0.9 weight_decay = 1e-4 epoch = 180 decay_milestones = [90, 120] decay_factor = 0.1 # SGM/L2 specific parameters aux_loss_wt = 0.02 train_transform = transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.RandomCrop(32, 4), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) test_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) image_datasets = { x: datasets.CIFAR10(root=args.cifar10_dir, train=y, download=True, transform=z) for x, y, z in zip([0, 1], [True, False], [train_transform, test_transform]) } dataloaders = { x: DataLoader(image_datasets[y], batch_size=batch_size, sampler=z, num_workers=args.num_workers, pin_memory=('cpu' not in args.device)) for x, y, z in zip(['train', 'val', 'test'], [0, 0, 1], [ sampler.SubsetRandomSampler(range(split_index)), sampler.SubsetRandomSampler( range(split_index, len(image_datasets[0]))), sampler.SequentialSampler(image_datasets[1]) ]) } dataset_sizes = { 'train': split_index, 'val': len(image_datasets[0]) - split_index, 'test': len(image_datasets[1]) } model = mutil.get_model(args.arch).to(device) if loss == 'default': criterion = torch.nn.CrossEntropyLoss().to(device) elif loss in ('sgm', 'l2'): criterion = GenericLoss(loss, aux_loss_wt, model.linear.out_features) else: raise NameError('{} is not recognized.'.format(loss)) optimizer = torch.optim.SGD(mutil.get_model_trainable_parameters(model), lr=lr, momentum=momentum, weight_decay=weight_decay) scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, milestones=decay_milestones, gamma=decay_factor) model, _ = mutil.train_model(model, criterion, optimizer, dataloaders, dataset_sizes, scheduler=scheduler, num_epochs=epoch, device=device) mutil.eval_model(model, dataloaders['test'], dataset_sizes['test'], device=device) return model
def run(args, ckpt_dir, ckpt_file): assert args.task == "snli" # Defining directories train_x, train_y, dev_x, dev_y, test_x, test_y, word_dict, embedding_matrix = load_all_data_snli(args) dev_matched_x, dev_matched_y, dev_mismatched_x, dev_mismatched_y = load_all_data_mnli(args, word_dict) vocab_size = embedding_matrix.shape[0] print("Dataset building all done") sess = tf.Session() use_additive = False if args.kwm_path != "": prev_arg_file = os.path.join(args.kwm_path, "args.pkl") prev_args = load_file(prev_arg_file) print("Loading key-word model with the following parameters: ") print(prev_args.__dict__) with tf.variable_scope(prev_args.modelname) as scope: prev_init = eval(model_utils.all_models[args.modeltype]) key_word_model = model_utils.get_model(prev_args, prev_init, vocab_size) kwm_saver = tf.train.Saver() kwm_ckpt = os.path.join(args.kwm_path, prev_args.modelname) kwm_saver.restore(sess, kwm_ckpt) use_additive = True with tf.variable_scope(args.modelname) as scope: init = eval(model_utils.all_models[args.modeltype]) pred_model = model_utils.get_model(args, init, vocab_size) saver = tf.train.Saver() if use_additive: init = models.AdditiveModel model = model_utils.get_additive_model(init, pred_model, key_word_model) else: model = pred_model utils.initialize_uninitialized_global_variables(sess) print("Building the model. Model name: {}".format(args.modelname)) if args.test: saver.restore(sess, ckpt_file) print('Test accuracy = ', model.evaluate_accuracy(sess, dev_x, dev_y)) else: sess.run(tf.assign(pred_model.embedding_w, embedding_matrix)) if os.path.exists(ckpt_file+".meta"): print('Restoring Model') saver.restore(sess, ckpt_file) print('Training..') for i in range(args.epochs): epoch_loss, epoch_accuracy = model.train_for_epoch(sess, train_x, train_y) print(i, 'loss: ', epoch_loss, 'acc: ', epoch_accuracy) # print('Train accuracy = ', model.evaluate_accuracy(sess, train_x, train_y)) # print(sess.run(tf.all_variables()[0][0])) print('Dev accuracy = ', model.evaluate_accuracy(sess, dev_x, dev_y)) print('Dev matched accuracy = ', model.evaluate_accuracy(sess, dev_matched_x, dev_matched_y)) print('Dev mismatched accuracy = ', model.evaluate_accuracy(sess, dev_mismatched_x, dev_mismatched_y)) if not os.path.exists(ckpt_dir): os.mkdir(ckpt_dir) print("Saving the model") saver.save(sess, ckpt_file) print("Finished") if model.use_alphas: print("Producing visualization") htmls = vis_utils.knit_nli(test_x, test_y, word_dict, None, model, sess, 100) f = open(os.path.join(ckpt_dir, "vis.html"), "wb") for i in htmls: f.write(i) f.close()
def train_model(config): if config.start_date is not None: print("Training start date: ", config.start_date) if config.start_date is not None: print("Training end date: ", config.end_date) print("Loading training data from %s ..." % config.datafile) train_data = None valid_data = None data_path = os.path.join(config.data_dir, config.datafile) batches = BatchGenerator(data_path, config, is_training_only=True) train_data = batches.train_batches(verbose=True) valid_data = batches.valid_batches(verbose=True) tf_config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False) with tf.Graph().as_default(), tf.Session(config=tf_config) as session: if config.seed is not None: tf.set_random_seed(config.seed) print("Constructing model ...") model = model_utils.get_model(session, config, verbose=True) params = model_utils.get_scaling_params(config, train_data, verbose=True) model.set_scaling_params(session, **params) noise_model = None if config.early_stop is not None: print("Training will early stop without " "improvement after %d epochs." % config.early_stop) sys.stdout.flush() train_history = list() valid_history = list() lr = model.set_learning_rate(session, config.learning_rate) train_data.cache(verbose=True) valid_data.cache(verbose=True) for i in range(config.max_epoch): (train_mse, valid_mse) = run_epoch(session, model, train_data, valid_data, keep_prob=config.keep_prob, passes=config.passes, noise_model=noise_model, verbose=True) print(( 'Epoch: %d Train MSE: %.6f Valid MSE: %.6f Learning rate: %.4f' ) % (i + 1, train_mse, valid_mse, lr)) sys.stdout.flush() train_history.append(train_mse) valid_history.append(valid_mse) if re.match("Gradient|Momentum", config.optimizer): lr = model_utils.adjust_learning_rate(session, model, lr, config.lr_decay, train_history) if not os.path.exists(config.model_dir): print("Creating directory %s" % config.model_dir) os.mkdir(config.model_dir) if math.isnan(valid_mse): print("Training failed due to nan.") quit() elif stop_training(config, valid_history): print("Training stopped.") quit() else: if ((config.early_stop is None) or (valid_history[-1] <= min(valid_history))): model_utils.save_model(session, config, i)
def verify(params): # get dataset ----------------------------------------------------------------------------------------------------- ans_to_id, id_to_ans = get_most_common_answers(vqa_train, vqa_val, int(params['num_answers']), []) embedding_matrix, ques_val_map, ans_val_map, img_val_map, ques_to_img_val = get_val_data([], False, True) normalize_image_embeddings([img_val_map]) nlp = spacy.load('en_vectors_glove_md') # get model ------------------------------------------------------------------------------------------------------- model = get_model( dropout_rate=0, regularization_rate=0, embedding_size=int(params['embedding_size']), num_classes=int(params['num_answers']), model_name=params['model'], embedding_matrix=embedding_matrix) savedir = "models/%s_%s" % (params['model'], str(params['num_answers'])) model.load_weights("%s/%s_epoch_%d_weights.h5" % (savedir, params['model'], params['epoch'])) layers = model.layers # get info about the layer weights--------------------------------------------------------------------------------- print "-------------------------------------------------" print "Layer weights:" print "-------------------------------------------------" for layer in layers: try: if layer.trainable: weights = layer.get_weights() print "Layer %s" % layer.name if len(weights) > 0: weights = weights[0] print "Shape {}".format(weights.shape) print "max %f" % np.max(weights) print "min %f" % np.min(weights) print "mean %f" % np.mean(weights) print "std %f" % np.std(weights) print "sum %f" % np.sum(weights) print "--------------------------" except: print "Layer %s - no weights" % layer.name print "--------------------------" pass # get some input image and question ------------------------------------------------------------------------------- vqa_ann_ids = vqa_val.getQuesIds() val_anns = vqa_val.loadQA(vqa_ann_ids) ann = val_anns[params['question_no']] question_prepro, image_prepro = ques_val_map[ann['question_id']], img_val_map[ann['image_id']] batch = [np.expand_dims(np.array(image_prepro), 0), np.expand_dims(np.array(question_prepro), 0)] pred = model.predict_on_batch(batch) pred_ans = np.argmax(pred[0]) print (id_to_ans[pred_ans]) print "-------------------------------------------------" print "Layer outputs:" print "-------------------------------------------------" for i in range(len(layers)): try: model_fn = K.function([layers[params['image_layer']].input, layers[params['text_layer']].input], [layers[i].output]) layer_output = model_fn(batch)[0] print "Layer %s" % layers[i].name print "Shape {}".format(layer_output.shape) print "max %f" % np.max(layer_output) print "min %f" % np.min(layer_output) print "mean %f" % np.mean(layer_output) print "std %f" % np.std(layer_output) print "sum %f" % np.sum(layer_output) print "--------------------------" except: print "Layer %s -- cannot show output" % layers[i].name print "--------------------------" pass fig = plt.figure() vqa_val.showQA([ann]) question = nlp(vqa_val.qqa[ann['question_id']]['question']) imgFilename = 'COCO_' + dataSubType_val + '_' + str(ann['image_id']).zfill(12) + '.jpg' if os.path.isfile(imgDir_val + imgFilename): I = io.imread(imgDir_val + imgFilename) I = scipy.misc.imresize(I, (2048, 2048), interp='nearest') fig.add_subplot(2, len(question), 1) plt.imshow(I) plt.xlabel(vqa_val.qqa[ann['question_id']]['question']) plt.xticks([]) plt.yticks([]) else: print(imgDir_val + imgFilename) print ("Multiple choice answer: %s" % ann['multiple_choice_answer']) model_fn = K.function([layers[params['image_layer']].input, layers[params['text_layer']].input], [layers[params['attention_layer']].output]) layer_output = model_fn(batch) layer_output = layer_output[0][0] layer_output_sum = np.sum(layer_output, axis=0) layer_output_img = scipy.misc.imresize(np.resize(layer_output_sum, (7, 7)), (2048, 2048), interp='bicubic') fig.add_subplot(2, len(question), 2) plot(I) plt.imshow(layer_output_img, cmap='gray', alpha=0.6) plt.xlabel("Ans: %s \nAns prob: %.2f" % (id_to_ans[pred_ans], np.max(pred))) plt.xticks([]) plt.yticks([]) for word_id in range(len(question) - 2): layer_output_img = scipy.misc.imresize(np.resize(layer_output[word_id], (7, 7)), (2048, 2048), interp='bicubic') print (layer_output[word_id], layer_output[word_id].shape) print "max %f" % np.max(layer_output[word_id]) print "min %f" % np.min(layer_output[word_id]) print "mean %f" % np.mean(layer_output[word_id]) print "std %f" % np.std(layer_output[word_id]) print "sum %f" % np.sum(layer_output[word_id]) fig.add_subplot(2, len(question), word_id + 3) plot(I) plt.imshow(layer_output_img, cmap='gray', alpha=0.6) plt.xlabel("%s %s %s" % (question[word_id].text, question[word_id + 1].text, question[word_id + 2].text)) plt.xticks([]) plt.yticks([]) model_fn = K.function([layers[params['image_layer']].input, layers[params['text_layer']].input], [layers[params['attention_layer'] + 1].output]) layer_output = model_fn(batch) layer_output = layer_output[0][0] layer_output_sum = np.sum(layer_output, axis=0) / len(layer_output) labels = ["%s\n%s\n%s\n" % (question[i].text, question[i+1].text, question[i+2].text) for i in range(len(question) - 2)] xlocations = np.array(range(len(layer_output_sum))) + 0.5 width = 0.5 fig.add_subplot(2, len(question) // 2, len(question) // 2 + 1) plt.bar(xlocations, layer_output_sum, width=width) plt.xticks(xlocations + width / 2, labels) plt.xlim(0, xlocations[-1] + width * 2) plt.show() # plt.savefig('attention_plot.png') plt.clf()
words_to_index, len(words_to_index), output_dim=config.dim) inp_type = tf.float32 print('Creating model') if config.sentence_embedding is not None: model = model_utils.get_SM_model_2(shape, embedding_layer=embeddings_layer, units=config.units, dtype=tf.string) else: model = model_utils.get_model(shape, embeddings_layer, config.classes, units=config.units, dtype=inp_type) model.summary() model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy']) callbacks = model_utils.get_callbacks( early_stop_monitor=config.early_stop_monitor, early_stop_patience=config.early_stop_patience, early_stop_mode=config.early_stop_mode) weights = model_utils.get_sample_weights_prim(train_y, config.class_weight) # model_info = model.fit(train_x_indices, train_y_oh, epochs=config.epochs, batch_size=config.batch_size, validation_split=0.05, callbacks=callbacks, shuffle=True, verbose=config.verbose)
def main(): global args, device args = parse_arguments() device = torch.device(args.device) imagenet32_labels = get_imagenet32_labels(args.imagenet32_dir) # Generate template network weights and last layer initialization weights template_file = os.path.join(args.model_dir, '{}.pth'.format(args.arch)) if args.overwrite or not os.path.isfile(template_file): print('Preparing {} template weights...'.format(args.arch)) model = train_template_network() pathlib.Path(os.path.dirname(template_file)).mkdir(parents=True, exist_ok=True) torch.save(model.state_dict(), template_file) else: print('{} template weights exist.'.format(args.arch)) if 'sgm' in args.experiments: sgm_file = os.path.join(args.model_dir, '{}_sgm.pth'.format(args.arch)) if args.overwrite or not os.path.isfile(sgm_file): print('Preparing {} SGM template weights...'.format(args.arch)) model = train_template_network(loss='sgm') pathlib.Path(os.path.dirname(sgm_file)).mkdir(parents=True, exist_ok=True) torch.save(model.state_dict(), sgm_file) else: print('{} SGM template weights exist.'.format(args.arch)) if 'l2' in args.experiments: l2_file = os.path.join(args.model_dir, '{}_l2.pth'.format(args.arch)) if args.overwrite or not os.path.isfile(l2_file): print('Preparing {} L2 template weights...'.format(args.arch)) model = train_template_network(loss='l2') pathlib.Path(os.path.dirname(l2_file)).mkdir(parents=True, exist_ok=True) torch.save(model.state_dict(), l2_file) else: print('{} L2 template weights exist.'.format(args.arch)) binary_linear_file = os.path.join(args.model_dir, '{}_binary_linear.pth'.format(args.arch)) if args.overwrite or not os.path.isfile(binary_linear_file): print('Preparing binary {} fully-connected weights...'.format( args.arch)) model = mutil.get_model(args.arch) linear = torch.nn.Linear(model.linear.in_features, 2) pathlib.Path(os.path.dirname(binary_linear_file)).mkdir(parents=True, exist_ok=True) torch.save(linear.state_dict(), binary_linear_file) else: print('Binary {} fully-connected weights exist.'.format(args.arch)) # Load component network class indices and experiment target indices. with open( os.path.join(args.indices_dir, 'imagenet_component_classes.json'), 'r') as f: component_classes = json.load(f) print('Index of {} component classes loaded.'.format( len(component_classes))) with open(os.path.join(args.indices_dir, 'imagenet_target_classes.json'), 'r') as f: target_classes = json.load(f) print('Index of {} target classes loaded.'.format(len(target_classes))) if any('combn' in x for x in args.experiments) or any('pcbn' in x for x in args.experiments): # Generate component networks (if haven't) for pos_class in component_classes: component_file = os.path.join( args.model_dir, COMPONENT_DIRNAME, '{}_{}.pth'.format(args.arch, pos_class)) if args.overwrite or not os.path.isfile(component_file): print('Training component network ({} {})...'.format( pos_class, imagenet32_labels[pos_class])) model = train_component_network( pos_class, template_file=template_file, binary_linear_file=binary_linear_file) pathlib.Path(os.path.dirname(component_file)).mkdir( parents=True, exist_ok=True) torch.save(model.state_dict(), component_file) # Evaluate component networks to rank them for selection, or load the # evaluations if they already exist if (any('accuracy' in x for x in args.experiments) or any('threshold' in x for x in args.experiments)): max_shot_eval_file = os.path.join(args.indices_dir, 'max_shot_accuracies.json') if args.overwrite or not os.path.isfile(max_shot_eval_file): print('Generating max-shot component accuracies...') all_accuracies = rank_component_networks(component_classes, target_classes, pos_size=0, method='accuracy') with open(max_shot_eval_file, 'w') as f: json.dump(all_accuracies, f) else: with open(max_shot_eval_file, 'r') as f: all_accuracies = json.load(f) print('Max-shot accuracy component evaluations loaded.') if any('loss' in x for x in args.experiments): few_shot_eval_file = os.path.join( args.indices_dir, '{}-shot_losses.json'.format(args.shot)) if args.overwrite or not os.path.isfile(few_shot_eval_file): print('Generating {}-shot component losses...'.format( args.shot)) all_losses = rank_component_networks(component_classes, target_classes, pos_size=args.shot, method='loss') with open(few_shot_eval_file, 'w') as f: json.dump(all_losses, f) else: with open(few_shot_eval_file, 'r') as f: all_losses = json.load(f) print('{}-shot loss component evaluations loaded.'.format( args.shot)) # Main experiment loop for experiment in args.experiments: shot_dir = 'max-shot' if args.shot == 0 else '{}-shot'.format( args.shot) if not args.train: # Perform evaluation by reading off the training summaries accuracies = [] for pos_class in target_classes: summary_file = os.path.join( args.model_dir, shot_dir, experiment, '{}_{}.summary'.format(args.arch, pos_class)) if os.path.isfile(summary_file): entry = torch.load(summary_file) accuracies.append(np.amax(entry['val_acc'])) if accuracies: print('Mean validation accuracy of {} ({} classes): {:.1f}%'. format(experiment, len(accuracies), np.mean(accuracies) * 100)) else: print( 'Mean validation accuracy of {} ({} classes): N/A'.format( experiment, len(accuracies))) else: for pos_class in target_classes: weights_file = os.path.join( args.model_dir, shot_dir, experiment, '{}_{}.pth'.format(args.arch, pos_class)) summary_file = os.path.join( args.model_dir, shot_dir, experiment, '{}_{}.summary'.format(args.arch, pos_class)) if not args.overwrite and os.path.isfile(weights_file): print('Weights found for {} ({} {}). Skipping...'.format( experiment, pos_class, imagenet32_labels[pos_class])) continue print('Preparing {} ({} {})...'.format( experiment, pos_class, imagenet32_labels[pos_class])) # Define model for this experiment if any(x in experiment for x in ('combn', 'pcbn')): # Parse experiment text to set up the proper BN combination # configuration exp_params = experiment.split('_') comb_method = exp_params[0] selection_params = {'method': exp_params[1]} if selection_params['method'] == 'loss': metrics = all_losses selection_params['num_components'] = int(exp_params[2]) elif selection_params['method'] == 'accuracy': metrics = all_accuracies selection_params['num_components'] = int(exp_params[2]) elif selection_params['method'] == 'threshold': metrics = all_accuracies selection_params['threshold'] = float(exp_params[2]) print('Selecting components...') comp_paths = select_components(pos_class, metrics, target_classes, component_classes, **selection_params) if comp_paths is None: print('No valid components. Skipping...') continue model = get_bn_combination_network( comp_paths, method=comb_method, template_file=template_file, binary_linear_file=binary_linear_file) elif experiment == 'last': model = mutil.get_model(args.arch) model.load_state_dict( torch.load(template_file, map_location='cpu')) mutil.freeze_model_parameters_(model) model.linear = torch.nn.Linear(model.linear.in_features, 2) model.linear.load_state_dict( torch.load(binary_linear_file, map_location='cpu')) elif experiment == 'full': model = mutil.get_model(args.arch) model.load_state_dict( torch.load(template_file, map_location='cpu')) model.linear = torch.nn.Linear(model.linear.in_features, 2) model.linear.load_state_dict( torch.load(binary_linear_file, map_location='cpu')) elif experiment == 'bn': model = mutil.get_model(args.arch) model.load_state_dict( torch.load(template_file, map_location='cpu')) mutil.freeze_model_parameters_(model) mutil.set_module_trainable_(model, torch.nn.BatchNorm2d) model.linear = torch.nn.Linear(model.linear.in_features, 2) model.linear.load_state_dict( torch.load(binary_linear_file, map_location='cpu')) elif experiment == 'sgm': model = mutil.get_model(args.arch) model.load_state_dict( torch.load(sgm_file, map_location='cpu')) mutil.freeze_model_parameters_(model) model.linear = torch.nn.Linear(model.linear.in_features, 2) model.linear.load_state_dict( torch.load(binary_linear_file, map_location='cpu')) elif experiment == 'l2': model = mutil.get_model(args.arch) model.load_state_dict( torch.load(l2_file, map_location='cpu')) mutil.freeze_model_parameters_(model) model.linear = torch.nn.Linear(model.linear.in_features, 2) model.linear.load_state_dict( torch.load(binary_linear_file, map_location='cpu')) else: raise NameError('{} is not recognized.'.format(experiment)) model.to(device) # Prepare dataset dataloaders = {} dataset_sizes = {} _, dataloaders['train'], dataset_sizes[ 'train'] = get_binary_imagenet32(pos_class, pos_size=args.shot, train=True) _, dataloaders['val'], dataset_sizes[ 'val'] = get_binary_imagenet32(pos_class, pos_size=0, train=False) # Train model and save weights optimizer = torch.optim.SGD( mutil.get_model_trainable_parameters(model), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, milestones=args.decay_milestones, gamma=args.decay_factor) print('Training...') model, summary = mutil.train_model( model, torch.nn.CrossEntropyLoss().to(device), optimizer, dataloaders, dataset_sizes, scheduler=scheduler, num_epochs=args.epoch, device=device, verbose=False) pathlib.Path(os.path.dirname(weights_file)).mkdir( parents=True, exist_ok=True) torch.save(model.state_dict(), weights_file) torch.save(summary, summary_file) print('Script complete.')
# print arguments print("\nParameters:") for attr, value in sorted(args.__dict__.items()): print("\t{}={}".format(attr.upper(), value)) # load vocabulary if args.snapshot is not None: vocab = pickle.load(open(args.snapshot + '.vocab', 'rb')) else: vocab = None # load data train_data_dict, dev_data_dict, test_data_dict, vocab = data_utils.load_dataset(args, vocab) # Load model model = model_utils.get_model(vocab, args) if args.mode == 'train_r2a': ''' Training R2A on labeled source and unlabeled target ''' dev_res, saved_path, model = train_utils.train(train_data_dict, dev_data_dict, model, args) # saving the vocabulary if args.save: with open(saved_path+'.vocab', 'wb') as f: pickle.dump(vocab, f, pickle.HIGHEST_PROTOCOL) # evaluate performance on the source train & dev set tar_train = None if args.tar_dataset == '' else train_data_dict[args.tar_dataset] tar_dev = None if args.tar_dataset == '' else dev_data_dict[args.tar_dataset]
parser.add_argument('--epoch_game_steps', type=int, default=10000, help='number of steps per epoch') parser.add_argument('--episode_discount', type=float, default=0.95, help='number of episodes for training') parser.add_argument('--seed', type=int, default=1, help='seed value') parser.add_argument( '--model', default='aac', choices=('aac', 'aac_lstm', 'aac_noisy', 'aac_depth', 'aac_map', 'mb_map'), help='model to work with') parser.add_argument('--base_model', default=None, help='path to base model file') parser.add_argument('--action_set', default=None, help='model to work with') parser.add_argument('--load', default=None, help='path to model file') parser.add_argument('--vizdoom_config', default='environments/basic.cfg', help='vizdoom config path') parser.add_argument('--vizdoom_path', default=_vzd_path, help='path to vizdoom') parser.add_argument('--wad_path', default=_vzd_path + '/freedoom2.wad', help='wad file path') parser.add_argument('--skiprate', type=int, default=1, help='number of skipped frames') parser.add_argument('--frame_num', type=int, default=1, help='number of frames per input') parser.add_argument('--checkpoint_file', default=None, help='check point file name') parser.add_argument('--checkpoint_rate', type=int, default=500, help='number of batches per checkpoit') parser.add_argument('--bot_cmd', default=None, help='command to launch a bot') parser.add_argument('--h5_path', default=None, help='hd5 files path') args = parser.parse_args() print(args) init_doom_env(args) model = get_model(args) if args.mode == 'train': model.run_train(args) else: model.run_test(args)
) display_predictions(convert_to_color(train_gt), viz, caption="Train ground truth") display_predictions(convert_to_color(test_gt), viz, caption="Test ground truth") # delete # display_predictions(convert_to_color(open_file('../Houston2013/gt.mat')['gt']), viz, caption="ground truth") if CLASS_BALANCING: weights = compute_imf_weights(train_gt, N_CLASSES, IGNORED_LABELS) hyperparams["weights"] = torch.from_numpy(weights) # Neural network model, optimizer, loss, hyperparams = get_model(MODEL, **hyperparams) # Split train set in train/val if SAMPLE_TRAIN_VALID != 1: train_gt, val_gt = sample_gt(train_gt, SAMPLE_TRAIN_VALID, mode="random") else: # Use all training data to train the model _, val_gt = sample_gt(train_gt, 0.95, mode="random") # Generate the dataset train_dataset = MultiModalX(img1, img2, train_gt, **hyperparams) train_loader = data.DataLoader( train_dataset, batch_size=hyperparams["batch_size"],
def __init__(self, config, model_file_path): self.config = config self.vocab = Vocab(config) self.model = get_model(self.vocab, config, model_file_path)
print('trial {} takes {:0.4f}secs: logit={:.5f} label={}'.format( i, dt, logits[top1], labelmap[top1])) print('average inference time {:.1f}ms'.format( exec_time / (n_trials - n_warmup) * 1e3)) # main entry if __name__ == '__main__': import argparse parser = argparse.ArgumentParser(description='model conversion test') parser.add_argument('--model_name', type=str, default='resnet50') parser.add_argument('--image_file', type=str, default='dog-komondor.jpg') args = parser.parse_args() # pytorch print('==== pytorch ====') model = get_model(args.model_name) n_params = count_param_size(model) print('{} has {:.4f}M trainable parameters'.format(args.model_name, n_params / 1e6)) target_size, trans = get_trans(args.model_name) flops, n_params = count_flops(model, target_size) print('THOP: {} has {:.4f}M trainable parameters'.format( args.model_name, n_params / 1e6)) print('THOP: {} has {:.4f}G FLOPs'.format(args.model_name, flops / 1e9)) logits_pytorch = run_inference_pytorch(args.image_file, model, trans) print_logits(logits_pytorch, args.model_name, 'pytorch') print('==== pytorch speed test ====') pytorch_speed_test(model, target_size)
def predict_multiple(images_batch, model, is_tensor=True, use_gpu=False): predictions = [] for image in images_batch: predictions.append(predict(image, model, is_tensor, use_gpu)) predictions = torch.cat(predictions) return predictions if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--arch', type=str, choices=ARCHS_LIST, default='resnet50') parser.add_argument('--image', type=str, required=True) args_dict = vars(parser.parse_args()) if os.path.exists(args_dict['image']): if args_dict['image'].endswith(('png', 'jpg', 'jpeg')): model = get_model(args_dict['arch'], parameters='standard').eval() predicted_class = torch.argmax( predict(args_dict['image'], model, is_tensor=False)).item() print(predicted_class) else: raise ValueError( 'The entered file is not an image with a format .png, .jpg or .jpeg!' ) else: raise ValueError('Incorrect image path!')