def train_angular(opt): logging.info(opt) # Settings. mx.random.seed(opt.seed) np.random.seed(opt.seed) # Setup computation context context = get_context(opt.gpus, logging) run_results = [] # Get model if opt.model == 'inception-bn': feature_net, feature_params = get_feature_model(opt.model, ctx=context) feature_net.collect_params().load(feature_params, ctx=context, ignore_extra=True) data_shape = 224 scale_image_data = False elif opt.model == 'resnet50_v2': feature_net = mx.gluon.model_zoo.vision.resnet50_v2( pretrained=True, ctx=context).features feature_params = mx.gluon.model_zoo.model_store.get_model_file( 'resnet%d_v%d' % (50, 2), root=os.path.join(mx.base.data_dir(), 'models')) data_shape = 224 scale_image_data = True else: raise RuntimeError('Unsupported model: %s' % opt.model) net = EmbeddingNet(feature_net, opt.embed_dim, normalize=False) if opt.model == 'resnet50_v2': # Use a smaller learning rate for pre-trained convolutional layers. for v in net.base_net.collect_params().values(): if 'conv' in v.name: setattr(v, 'lr_mult', 0.01) # Get iterators train_data, _ = get_npairs_iterators( opt.dataset, opt.data_path, batch_size=opt.batch_size, data_shape=data_shape, test_batch_size=len(context) * 32, use_crops=opt.use_crops, scale_image_data=scale_image_data, same_image_sampling=opt.same_image_sampling) train_it_dataloader = mx.gluon.data.DataLoader(DatasetIterator( train_data, opt.epoch_length), batch_size=1, shuffle=False, num_workers=opt.num_workers, last_batch='rollover') _, val_dataset = get_dataset(opt.dataset, opt.data_path, data_shape=data_shape, use_crops=opt.use_crops, use_aug=True, with_proxy=False, scale_image_data=scale_image_data) val_data = mx.gluon.data.DataLoader(val_dataset, batch_size=opt.batch_size, shuffle=False, num_workers=opt.num_workers, last_batch='keep') for run in range(1, opt.number_of_runs + 1): logging.info('Starting run %d/%d' % (run, opt.number_of_runs)) net.init(mx.init.Xavier(magnitude=0.2), ctx=context, init_basenet=False) if opt.model == 'inception-bn': net.base_net.collect_params().load(feature_params, ctx=context, ignore_extra=True) elif opt.model == 'resnet50_v2': net.base_net.load_parameters(feature_params, ctx=context, allow_missing=True, ignore_extra=True) if not opt.disable_hybridize: net.hybridize() run_result = train(net, opt, train_it_dataloader, val_data, context, run) run_results.append(run_result) logging.info('Run %d finished with %f' % (run, run_result[0][1])) logging.info( 'Average validation of %d runs:\n%s' % (opt.number_of_runs, format_results(average_results(run_results))))
def train_discriminative(opt): logging.info(opt) # Settings. mx.random.seed(opt.seed) np.random.seed(opt.seed) # Setup computation context context = get_context(opt.gpus, logging) run_results = [] # Get model if opt.model == 'inception-bn': feature_net, feature_params = get_feature_model(opt.model, ctx=context) feature_net.collect_params().load(feature_params, ctx=context, ignore_extra=True) data_shape = 224 scale_image_data = False elif opt.model == 'resnet50_v2': feature_net = mx.gluon.model_zoo.vision.resnet50_v2( pretrained=True, ctx=context).features data_shape = 224 scale_image_data = True else: raise RuntimeError('Unsupported model: %s' % opt.model) # Get data iterators train_dataset, val_dataset = get_dataset(opt.dataset, opt.data_path, data_shape=data_shape, use_crops=opt.use_crops, use_aug=True, with_proxy=True, scale_image_data=scale_image_data) logging.info('Training with %d classes, validating with %d classes' % (train_dataset.num_classes(), val_dataset.num_classes())) train_dataloader = mx.gluon.data.DataLoader(train_dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers, last_batch='rollover') val_dataloader = mx.gluon.data.DataLoader(val_dataset, batch_size=opt.batch_size, shuffle=False, num_workers=opt.num_workers, last_batch='keep') net = EmbeddingNet( feature_net, [opt.embed_dim, train_dataset.num_classes()], normalize=True, dropout=True) if opt.model == 'resnet50_v2': # Use a smaller learning rate for pre-trained convolutional layers. for v in net.base_net.collect_params().values(): if 'conv' in v.name: setattr(v, 'lr_mult', 0.01) # main run loop for multiple training runs for run in range(1, opt.number_of_runs + 1): logging.info('Starting run %d/%d' % (run, opt.number_of_runs)) net.init(mx.init.Xavier(magnitude=0.2), ctx=context, init_basenet=False) if opt.model == 'inception-bn': net.base_net.collect_params().load(feature_params, ctx=context, ignore_extra=True) elif opt.model == 'resnet50_v2': net.base_net = mx.gluon.model_zoo.vision.resnet50_v2( pretrained=True, ctx=context).features else: raise RuntimeError('Unsupported model: %s' % opt.model) if not opt.disable_hybridize: net.hybridize() run_result = train(net, opt, train_dataloader, val_dataloader, context, run) run_results.append(run_result) logging.info('Run %d finished with %f' % (run, run_result[0][1])) logging.info( 'Average validation of %d runs:\n%s' % (opt.number_of_runs, format_results(average_results(run_results))))
def train_rankedlist(opt): logging.info(opt) # Settings. mx.random.seed(opt.seed) np.random.seed(opt.seed) # Setup computation context context = get_context(opt.gpus, logging) run_results = [] # Get model if opt.model == 'inception-bn': feature_net, feature_params = get_feature_model(opt.model, ctx=context) feature_net.collect_params().load(feature_params, ctx=context, ignore_extra=True) data_shape = 224 scale_image_data = False elif opt.model == 'resnet50_v2': feature_net = mx.gluon.model_zoo.vision.resnet50_v2( pretrained=True, ctx=context).features data_shape = 224 scale_image_data = True feature_params = None else: raise RuntimeError('Unsupported model: %s' % opt.model) if opt.bottleneck_layers != '': embedding_layers = [int(x) for x in opt.bottleneck_layers.split(',') ] + [opt.embed_dim] else: embedding_layers = [opt.embed_dim] logging.info('Embedding layers: [%s]' % ','.join([str(x) for x in embedding_layers])) if len(embedding_layers) == 1: embedding_layers = embedding_layers[0] net = EmbeddingNet(feature_net, embedding_layers, normalize=True, dropout=False) logging.info(net) if opt.model == 'resnet50_v2': # Use a smaller learning rate for pre-trained convolutional layers. for v in net.base_net.collect_params().values(): if 'conv' in v.name: setattr(v, 'lr_mult', 0.01) elif 'batchnorm' in v.name or 'bn_' in v.name: v.grad_req = 'null' else: for v in net.encoder.collect_params().values(): setattr(v, 'lr_mult', 10.) # Get data iterators train_dataset = DatasetIterator( get_dataset_iterator(opt.dataset, opt.data_path, batch_k=opt.batch_k, batch_size=opt.batch_size, batchify=False, data_shape=data_shape, use_crops=opt.use_crops, scale_image_data=scale_image_data)[0], opt.iteration_per_epoch, 'next') train_dataiterator = mx.gluon.data.DataLoader(train_dataset, batch_size=1, shuffle=False, num_workers=opt.num_workers, last_batch='keep') val_dataset = get_dataset(opt.dataset, opt.data_path, data_shape=data_shape, use_crops=opt.use_crops, use_aug=True, scale_image_data=scale_image_data)[1] logging.info( 'Training with %d classes, validating with %d classes' % (train_dataset.data_iterator.num_classes(), val_dataset.num_classes())) val_dataloader = mx.gluon.data.DataLoader(val_dataset, batch_size=opt.batch_size, shuffle=False, num_workers=opt.num_workers, last_batch='keep') # main run loop for multiple training runs for run in range(1, opt.number_of_runs + 1): logging.info('Starting run %d/%d' % (run, opt.number_of_runs)) net.init(mx.init.Xavier(magnitude=0.2), ctx=context, init_basenet=False) if opt.model == 'inception-bn': net.base_net.collect_params().load(feature_params, ctx=context, ignore_extra=True) elif opt.model == 'resnet50_v2': net.base_net = mx.gluon.model_zoo.vision.resnet50_v2( pretrained=True, ctx=context).features else: raise RuntimeError('Unsupported model: %s' % opt.model) if not opt.disable_hybridize: net.hybridize() run_result = train(net, opt, train_dataiterator, val_dataloader, context, run) run_results.append(run_result) logging.info('Run %d finished with %f' % (run, run_result[0][1])) logging.info( 'Average validation of %d runs:\n%s' % (opt.number_of_runs, format_results(average_results(run_results))))