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