def build_training_model(config, experiment):
    # set device
    device_opt = caffe2_pb2.DeviceOption()
    if config['gpu_id'] is not None:
        device_opt.device_type = caffe2_pb2.CUDA
        device_opt.cuda_gpu_id = config['gpu_id']

    # build model
    with core.DeviceScope(device_opt):
        training_model = model_helper.ModelHelper(
            name='{}_training_model'.format(config['name']), )
        data, label = add_input(training_model, config, is_test=False)
        pred = add_model(training_model, config, data, is_test=False)
        loss = add_loss(training_model, config, pred, label)
        add_training_operators(training_model, config, loss)
        add_accuracy(training_model)

    # init workspace for training net
    workspace.RunNetOnce(training_model.param_init_net)

    # if in finetune mode, we need to load pretrained weights and bias
    if config['finetune']:
        load_init_net(config['network']['init_net'], device_opt)
        if config['solver']['percent'] is not None:
            prune_init_net(workspace, config, device_opt, experiment)

    workspace.CreateNet(training_model.net)
    return training_model
Exemplo n.º 2
0
def build_validation_model(config):
    # set device
    device_opt = caffe2_pb2.DeviceOption()
    if config['gpu_id'] is not None:
        device_opt.device_type = caffe2_pb2.CUDA
        device_opt.cuda_gpu_id = config['gpu_id']

    # build model
    with core.DeviceScope(device_opt):
        validation_model = model_helper.ModelHelper(
            name='{}_validation_model'.format(config['name']),
            init_params=False,
        )
        data, label = add_input(validation_model, config, is_test=True)
        pred = add_model(validation_model, config, data, is_test=True)
        loss = add_loss(validation_model, config, pred, label)
        add_accuracy(validation_model)

    # init workspace for validation net
    workspace.RunNetOnce(validation_model.param_init_net)

    # load pretrained network params
    load_init_net(config['network']['init_net'], device_opt)

    workspace.CreateNet(validation_model.net)
    return validation_model