def createTestModel(lmdb_path):
"""Create and return a test model. Does not include training ops."""
model = ModelHelper(name='test',
arg_scope={'order': 'NCHW'},
init_params=False)
reader = model.CreateDB('test_reader', db=lmdb_path, db_type='lmdb')
AddInputOps(model, reader, BATCH_SIZE)
AddForwardPassOps(model)
workspace.RunNetOnce(model.param_init_net)
workspace.CreateNet(model.net)
return model
def createTrainModel(lmdb_path):
"""Create and return a training model, complete with training ops."""
model = ModelHelper(name='train', arg_scope={'order': 'NCHW'})
reader = model.CreateDB('train_reader', db=lmdb_path, db_type='lmdb')
AddInputOps(model, reader, BATCH_SIZE)
losses = AddForwardPassOps(model)
model.AddGradientOperators(losses)
AddOptimizerOps(model)
workspace.RunNetOnce(model.param_init_net)
workspace.CreateNet(model.net)
return model
def createTrainModel(lmdb_path, devices):
"""Create and return a training model, complete with training ops."""
model = ModelHelper(name='train', arg_scope={'order': 'NCHW'})
reader = model.CreateDB('train_reader', db=lmdb_path, db_type='lmdb')
data_parallel_model.Parallelize_GPU(
model,
input_builder_fun=functools.partial(AddInputOps,
reader=reader,
batch_size=(BATCH_SIZE //
len(devices))),
forward_pass_builder_fun=AddForwardPassOps,
optimizer_builder_fun=AddOptimizerOps,
devices=devices,
use_nccl=True)
workspace.RunNetOnce(model.param_init_net)
workspace.CreateNet(model.net)
return model
def createTestModel(lmdb_path, devices):
"""Create and return a test model. Does not include training ops."""
model = ModelHelper(name='test',
arg_scope={'order': 'NCHW'},
init_params=False)
reader = model.CreateDB('test_reader', db=lmdb_path, db_type='lmdb')
data_parallel_model.Parallelize_GPU(
model,
input_builder_fun=functools.partial(AddInputOps,
reader=reader,
batch_size=(BATCH_SIZE //
len(devices))),
forward_pass_builder_fun=AddForwardPassOps,
param_update_builder_fun=None,
devices=devices)
workspace.RunNetOnce(model.param_init_net)
workspace.CreateNet(model.net)
return model
