def main(): opt = parse_args() use_cuda = cuda_mode(opt) save_dir = os.path.join(opt.out_dir, opt.dataset_name) print('Trained models will be save to', os.path.abspath(save_dir)) if not os.path.exists(save_dir): os.makedirs(save_dir) pre_model = opt.pre_model if opt.resume != 0: pre_model = '' model = SIPN(opt.net, opt.dataset_name, pre_model) if opt.resume: resume = os.path.join(save_dir, 'sipn_' + opt.net + '_' + str(opt.resume) + '.pth') print('Resuming model check point from {}'.format(resume)) model.load_trained_model(torch.load(resume)) # Load the dataset dataset = PersonSearchDataset(opt.data_dir, opt.dataset_name) # Choose parameters to be updated during training lr = opt.lr params = [] print('These parameters will be updated during training:') for key, value in dict(model.named_parameters()).items(): if value.requires_grad: print(key) # TODO: set different decay for weight and bias params += [{'params': [value], 'lr': lr, 'weight_decay': 1e-4}] if opt.optimizer == 'SGD': optimizer = torch.optim.SGD(params, momentum=0.9) elif opt.optimizer == 'Adam': # lr *= 0.1 optimizer = torch.optim.Adam(params) else: raise KeyError(opt.optimizer) # Train the model train_model(dataset, model, lr, optimizer, opt.epochs, use_cuda, save_dir, opt.resume) print('Done.\n')
def main(): """Test the model""" opt = parse_args() test_result_dir = os.path.join(opt.out_dir, 'test_result') dataset_gallery = PersonSearchDataset(opt.data_dir, split_name='test', gallery_size=200) if opt.use_saved_result: gboxes_file = open(os.path.join(test_result_dir, 'gboxes.pkl'), 'rb') g_boxes = pickle.load(gboxes_file) gfeatures_file = open(os.path.join(test_result_dir, 'gfeatures.pkl'), 'rb') g_features = pickle.load(gfeatures_file) qfeatures_file = open(os.path.join(test_result_dir, 'qfeatures.pkl'), 'rb') q_features = pickle.load(qfeatures_file) else: use_cuda = cuda_mode(opt) trained_model_dir = os.path.join('./output', 'sipn_' + opt.trained_epochs + '.pth') if not os.path.exists(test_result_dir): os.makedirs(test_result_dir) net = SIPN(opt.net, trained_model_dir, is_train=False) net.eval() if use_cuda: net.cuda() # load trained model print('Loading model check point from {:s}'.format(trained_model_dir)) net.load_trained_model(torch.load(trained_model_dir)) dataset_query = PersonSearchDataset(opt.data_dir, split_name='test', gallery_size=200, test_mode='query') g_boxes, g_features = test_gallery(net, dataset_gallery, use_cuda, test_result_dir) q_features = test_query(net, dataset_query, use_cuda, test_result_dir) # dataset_gallery.evaluate_detections(g_boxes, det_thresh=0.5) dataset_gallery.evaluate_search(g_boxes, g_features, q_features, det_thresh=0.5, gallery_size=200, dump_json=None)
def main(): opt = parse_args() use_cuda = cuda_mode(opt) model = SIPN(opt.net, opt.pre_model) # Load the dataset dataset = PersonSearchDataset(opt.data_dir) save_dir = opt.out_dir print('Trained models will be save to', os.path.abspath(save_dir)) if not os.path.exists(save_dir): os.makedirs(save_dir) # Choose parameters to be updated during training lr = opt.lr params = [] print('These parameters will be updated during training:') for key, value in dict(model.named_parameters()).items(): if value.requires_grad: print(key) # TODO: set different decay for weight and bias params += [{'params': [value], 'lr': lr, 'weight_decay': 1e-4}] if opt.optimizer == 'SGD': optimizer = torch.optim.SGD(params, momentum=0.9) elif opt.optimizer == 'Adam': lr *= 0.1 optimizer = torch.optim.Adam(params) else: raise KeyError(opt.optimizer) # TODO: add resume # Train the model train_model(dataset, model, lr, optimizer, opt.epochs, use_cuda, save_dir) print('Done')
def main(): opt = parse_args() use_cuda = cuda_mode(opt) trained_model_dir = os.path.join('./output', 'sipn_' + opt.net + '_' + opt.trained_epochs + '.pth') net = SIPN(opt.net, trained_model_dir, is_train=False) net.eval() if use_cuda: net.cuda() # load trained model print('Loading model check point from {:s}'.format(trained_model_dir)) net.load_trained_model(torch.load(trained_model_dir)) test_images = os.listdir(opt.data_dir) # demo_detection(net, opt.data_dir, test_images, use_cuda) demo_search(net, opt.data_dir, test_images, use_cuda)
def main(): opt = parse_args() global device device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') save_dir = os.path.join(opt.out_dir, opt.dataset_name) print('Trained models will be save to', os.path.abspath(save_dir)) if not os.path.exists(save_dir): os.makedirs(save_dir) pre_model = opt.pre_model if opt.resume != 0: pre_model = '' model = SIPN(opt.net, opt.dataset_name, pre_model) model.to(device) if opt.resume: resume = os.path.join(save_dir, 'sipn_' + opt.net + '_' + str(opt.resume) + '.pth') print('Resuming model check point from {}'.format(resume)) model.load_trained_model(torch.load(resume)) # Read the configuration file with open('config.yml', 'r') as f: config = yaml.load(f) target_size = config['target_size'] max_size = config['max_size'] pixel_means = config['pixel_means'] # Compose transformations for the dataset transform = sipn_transforms.Compose([ sipn_transforms.RandomHorizontalFlip(), sipn_transforms.Scale(target_size, max_size), sipn_transforms.ToTensor(), sipn_transforms.Normalize(pixel_means) ]) # Load the dataset dataset = SIPNDataset(opt.data_dir, opt.dataset_name, 'train', transform) dataloader = DataLoader(dataset, shuffle=True, num_workers=8) # Choose parameters to be updated during training lr = opt.lr params = [] print('These parameters will be updated during training:') for key, value in dict(model.named_parameters()).items(): if value.requires_grad: print(key) # TODO: set different decay for weight and bias params += [{'params': [value], 'lr': lr, 'weight_decay': 1e-4}] if opt.optimizer == 'SGD': optimizer = torch.optim.SGD(params, momentum=0.9) elif opt.optimizer == 'Adam': # lr *= 0.1 optimizer = torch.optim.Adam(params) else: raise KeyError(opt.optimizer) # Train the model train_model(dataloader, model, lr, optimizer, opt.epochs, save_dir, opt.resume) print('Done.\n')
color='white') plt.tight_layout() fig.savefig(os.path.join(im_dir, 'result_' + im_name)) plt.show() plt.close(fig) if __name__ == '__main__': opt = parse_args() use_cuda = cuda_mode(opt) trained_model_dir = os.path.join('./output', 'sipn_' + opt.trained_epochs + '.pth') net = SIPN(opt.net, trained_model_dir, training=False) net.eval() if use_cuda: net.cuda() # load trained model print('Loading model check point from {:s}'.format(trained_model_dir)) net.load_trained_model(torch.load(trained_model_dir)) test_images = os.listdir(opt.data_dir) # demo_detection(net, opt.data_dir, test_images, use_cuda) demo_search(net, opt.data_dir, test_images, use_cuda) print('Done')