示例#1
0
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))))
示例#3
0
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))))