def __init__(self, io_cfg, chain_cfg, verbose=False):
'''
Initializes the chain from the configuration file
'''
# Initialize the data loader
io_cfg = yaml.load(io_cfg, Loader=yaml.Loader)
# Save config, initialize output
self.cfg = chain_cfg
self.verbose = verbose
self.output = {}
# Initialize log
log_path = chain_cfg['name'] + '_log.csv'
print('Initialized Pi0 mass chain, log path:', log_path)
self._log = CSVData(log_path)
self._keys = ['event_id', 'pion_id', 'pion_mass']
# If a network is specified, initialize the network
self.network = False
if chain_cfg['segment'] == 'uresnet' or chain_cfg[
'shower_start'] == 'ppn':
self.network = True
with open(chain_cfg['net_cfg']) as cfg_file:
net_cfg = yaml.load(cfg_file, Loader=yaml.Loader)
io_cfg['model'] = net_cfg['model']
io_cfg['trainval'] = net_cfg['trainval']
# Pre-process configuration
process_config(io_cfg)
# Instantiate "handlers" (IO tools)
self.hs = prepare(io_cfg)
self.data_set = iter(self.hs.data_io)
def main():
cfg_file = sys.argv[1]
if not os.path.isfile(cfg_file):
cfg_file = os.path.join(current_directory, 'config', sys.argv[1])
if not os.path.isfile(cfg_file):
print(sys.argv[1], 'not found...')
sys.exit(1)
ckpt_dir = sys.argv[2]
if not os.path.isdir(ckpt_dir):
print(sys.argv[2], ' not a valid directory!')
sys.exit(1)
cfg = yaml.load(open(cfg_file, 'r'), Loader=yaml.Loader)
process_config(cfg)
log_dir = cfg['training']['log_dir']
# loop over configuration files
wfiles = np.array([w for w in os.listdir(ckpt_dir) if w.endswith('.ckpt')])
wfiles = np.sort(wfiles)
for wfile in wfiles:
print(wfile)
filename, file_extension = os.path.splitext(wfile)
# set weight file
cfg['training']['model_path'] = ckpt_dir + '/' + wfile
# set output log
cfg['training']['log_dir'] = log_dir + '/' + filename
# run inference
inference(cfg)
def __init__(self, io_cfg, chain_cfg, verbose=False):
'''
Initializes the chain from the configuration file
'''
# Initialize the data loader
io_cfg = yaml.load(io_cfg, Loader=yaml.Loader)
# Save config, initialize output
self.cfg = chain_cfg
self.verbose = verbose
self.event = None
self.output = {}
# Initialize log
log_path = chain_cfg['name'] + '_log.csv'
print('Initialized Pi0 mass chain, log path:', log_path)
self._log = CSVData(log_path)
self._keys = ['event_id', 'pion_id', 'pion_mass']
# If a network is specified, initialize the network
self.network = False
if chain_cfg['segment'] == 'uresnet' or chain_cfg[
'shower_start'] == 'ppn':
self.network = True
with open(chain_cfg['net_cfg']) as cfg_file:
net_cfg = yaml.load(cfg_file, Loader=yaml.Loader)
io_cfg['model'] = net_cfg['model']
io_cfg['trainval'] = net_cfg['trainval']
# Initialize the fragment identifier
self.frag_est = FragmentEstimator()
# If a direction estimator is requested, initialize it
if chain_cfg['shower_dir'] != 'truth':
self.dir_est = DirectionEstimator()
# If a clusterer is requested, initialize it
if chain_cfg['shower_energy'] == 'cone':
self.clusterer = ConeClusterer()
# If a pi0 identifier is requested, initialize it
if chain_cfg['shower_match'] == 'proximity':
self.matcher = Pi0Matcher()
# Pre-process configuration
process_config(io_cfg)
# Instantiate "handlers" (IO tools)
self.hs = prepare(io_cfg)
self.data_set = iter(self.hs.data_io)
def __init__(self, io_cfg, chain_cfg, verbose=False):
'''
Initializes the chain from the configuration file
'''
# Initialize the data loader
io_cfg = yaml.load(io_cfg, Loader=yaml.Loader)
# Save config, initialize output
self.cfg = chain_cfg
self.verbose = verbose
self.event = None
self.output = {}
# Initialize log
print('Initialized Pi0 reconstruction chain')
self._log = Pi0DataLogger(chain_cfg['name'])
# If a network is specified, initialize the network
self.network = False
if chain_cfg['segment'] == 'uresnet' or chain_cfg[
'shower_start'] == 'ppn' or chain_cfg['shower_start'] == 'gnn':
self.network = True
with open(chain_cfg['net_cfg']) as cfg_file:
net_cfg = yaml.load(cfg_file, Loader=yaml.Loader)
io_cfg['model'] = net_cfg['model']
io_cfg['trainval'] = net_cfg['trainval']
# Initialize the fragment identifier
if chain_cfg['shower_fragment'] == 'dbscan':
self.frag_est = DBSCANCluster()
# If a direction estimator is requested, initialize it
if chain_cfg['shower_direction'] != 'label':
self.dir_est = DirectionEstimator()
# If a clusterer is requested, initialize it
if chain_cfg['shower_cluster'] in ['cone', 'gnn']:
self.clusterer = ConeClusterer()
# If a pi0 identifier is requested, initialize it
if chain_cfg['shower_match'] == 'proximity':
self.matcher = Pi0Matcher()
# Pre-process configuration
process_config(io_cfg)
# Instantiate "handlers" (IO tools)
self.hs = prepare(io_cfg)
self.data_set = iter(self.hs.data_io)
def main():
cfg_file = sys.argv[1]
if not os.path.isfile(cfg_file):
cfg_file = os.path.join(current_directory, 'config', sys.argv[1])
if not os.path.isfile(cfg_file):
print(sys.argv[1], 'not found...')
sys.exit(1)
cfg = yaml.load(open(cfg_file, 'r'), Loader=yaml.Loader)
process_config(cfg)
if cfg['training']['train']:
train(cfg)
else:
inference(cfg)
def test_model_full(config_full):
"""
Tests whether a model can be trained.
Including parsers and trainval in the execution.
Parameters
----------
config: dict
Generated by a fixture above, dummy config to allow networks to run.
It is mostly empty, we rely on networks default config.
"""
config = config_full
model, criterion = factories.construct(config['model']['name'])
net = model(config['model'])
loss = criterion(config['model'])
if not hasattr(net, "INPUT_SCHEMA"):
pytest.skip('No test defined for network of %s' %
config['model']['name'])
if not hasattr(loss, "INPUT_SCHEMA"):
pytest.skip('No test defined for criterion of %s' %
config['model']['name'])
# Setup configuration to have all necessary I/O keys
config['iotool']['dataset']['schema'] = {}
config['model']['network_input'] = []
config['model']['loss_input'] = []
for i, x in enumerate(net.INPUT_SCHEMA + loss.INPUT_SCHEMA):
parser_name = x[0]
parser_return_types = x[1]
config['iotool']['dataset']['schema'][i] = [x[0]]
for t in parser_return_types:
config['iotool']['dataset']['schema'][i].extend(
branch[parser_name][t])
if i < len(net.INPUT_SCHEMA):
config['model']['network_input'].append(i)
else:
config['model']['loss_input'].append(i)
process_config(config)
# Try opening LArCV data file
try:
handlers = prepare(config)
except FileNotFoundError:
pytest.skip('File not found to test the loader.')
train_loop(config, handlers)
def main():
cfg_file = sys.argv[1]
if not os.path.isfile(cfg_file):
cfg_file = os.path.join(current_directory, 'config', sys.argv[1])
if not os.path.isfile(cfg_file):
print(sys.argv[1], 'not found...')
sys.exit(1)
cfg = yaml.load(open(cfg_file, 'r'), Loader=yaml.Loader)
if environ.get('CUDA_VISIBLE_DEVICES') is not None and cfg['trainval']['gpus'] == '-1':
cfg['trainval']['gpus'] = os.getenv('CUDA_VISIBLE_DEVICES')
process_config(cfg)
if cfg['trainval']['train']:
train(cfg)
else:
inference(cfg)
seed: 0
learning_rate: 0.0025
gpus: '0'
weight_prefix: '/gpfs/slac/staas/fs1/g/neutrino/qatom/gnn-models/early_stop_pi0-20.ckpt'
iterations: 1000
report_step: 1
checkpoint_step: 100
log_dir: logs/edge_gnn/edge_node_only
model_path: ''
train: True
debug: False
minibatch_size: 1
'''
cfg = yaml.load(cfg, Loader=yaml.Loader)
process_config(cfg)
loader, cfg['data_keys'] = loader_factory(cfg)
import torch
import pickle
from train_voxel_gnn import GraphDataset
from os import listdir
from torch.utils.data import DataLoader
from mlreco.main_funcs import cycle
data_path = '/gpfs/slac/staas/fs1/g/neutrino/qatom/gnn_pi0_reco_nocompton/'
valid_data = GraphDataset(data_path, 1, data_key='00000')
valid_dataset = cycle(valid_data)
from mlreco.trainval import trainval