def prepare(cfg):
torch.cuda.set_device(cfg['training']['gpus'][0])
handlers = Handlers()
# IO configuration
# Batch size for I/O becomes minibatch size
batch_size = cfg['iotool']['batch_size']
cfg['iotool']['batch_size'] = cfg['training']['minibatch_size']
handlers.data_io, cfg['data_keys'] = loader_factory(cfg)
# TODO check that it does what we want (cycle through dataloader)
# check on a small sample, check 1/ it cycles through and 2/ randomness
if cfg['training']['train']:
handlers.data_io_iter = iter(cycle(handlers.data_io))
else:
handlers.data_io_iter = itertools.cycle(handlers.data_io)
cfg['iotool']['batch_size'] = batch_size
# Trainer configuration
handlers.trainer = trainval(cfg)
# Restore weights if necessary
loaded_iteration = handlers.trainer.initialize()
if cfg['training']['train']:
handlers.iteration = loaded_iteration
make_directories(cfg, loaded_iteration, handlers=handlers)
return handlers
def test_model_train(config):
"""
TODO should test whether a model can be trained.
Need to write a fixture to generate dummy input data. Maybe we need to
explicitly define 1 test / model instead of a fixture?
"""
trainer = trainval(config)
def prepare(cfg):
"""
Prepares high level API handlers, namely trainval instance and torch DataLoader (w/ iterator)
INPUT
- cfg is a full configuration block after pre-processed by process_config function
OUTPUT
- Handler instance attached with trainval/DataLoader instances (if in config)
"""
handlers = Handlers()
handlers.cfg = cfg
# Instantiate DataLoader
handlers.data_io = loader_factory(cfg)
# IO iterator
handlers.data_io_iter = itertools.cycle(handlers.data_io)
if 'trainval' in cfg:
# Set random seed for reproducibility
np.random.seed(cfg['trainval']['seed'])
torch.manual_seed(cfg['trainval']['seed'])
# Set primary device
if len(cfg['trainval']['gpus']) > 0:
torch.cuda.set_device(cfg['trainval']['gpus'][0])
# TODO check that it does what we want (cycle through dataloader)
# check on a small sample, check 1/ it cycles through and 2/ randomness
if cfg['trainval']['train']:
handlers.data_io_iter = iter(cycle(handlers.data_io))
# Trainer configuration
handlers.trainer = trainval(cfg)
# set the shared clock
handlers.watch = handlers.trainer._watch
# Restore weights if necessary
loaded_iteration = handlers.trainer.initialize()
if cfg['trainval']['train']:
handlers.iteration = loaded_iteration
make_directories(cfg, loaded_iteration, handlers=handlers)
return handlers
valid_dataset = cycle(valid_data)
from mlreco.trainval import trainval
from mlreco.main_funcs import get_data_minibatched
nf_dim, ef_dim = valid_data.get_feature_dimensions()
cfg['model']['modules']['edge_model']['model_cfg']['n_node_features'] = nf_dim
cfg['model']['modules']['edge_model']['model_cfg']['n_edge_features'] = ef_dim
cfg['model']['network_input'] = [
'vertex_features', 'edge_features', 'edge_index', 'batch'
]
cfg['model']['loss_input'] = ['target', 'edge_index', 'batch']
cfg['data_keys'] = [
'edge_index', 'vertex_features', 'edge_features', 'batch', 'target'
]
Trainer = trainval(cfg)
loaded_iteration = Trainer.initialize()
from mlreco.utils.gnn.features.utils import edge_labels_to_node_labels
from mlreco.utils.metrics import *
valid_dataset = cycle(valid_data)
pred_sbd = []
for t in range(len(valid_data)):
data_blob = get_data_minibatched(valid_dataset, cfg)
Trainer._train = False
res = Trainer.forward(data_blob)
data_blob = get_data_minibatched(valid_dataset, cfg)
input_keys = Trainer._input_keys
loss_keys = Trainer._loss_keys
target = data_blob['target'][0][0]