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