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))))