This notebook is a basic demo of using lartpc_mlreco3d API to train and run inference, as well as display the events and the network output. We will focus on simple UResNet network using SCN library.
You can start by cloning the current code from Github repository.
Note: Use the develop branch.
import numpy as np
import matplotlib
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import pyplot as plt
import pandas as pd
import seaborn
seaborn.set(rc={
'figure.figsize':(15, 10),
})
seaborn.set_context('talk') # or paper
import plotly.plotly as py
import plotly.graph_objs as go
from plotly.offline import download_plotlyjs, init_notebook_mode, plot, iplot
init_notebook_mode(connected=False)
Let's load the lartpc_mlreco3d module first (change the path to wherever you cloned the repository):
import sys
sys.path.append("/u/ki/ldomine/lartpc_mlreco3d/")
# If necessary reload the mlreco module after modifying it
# import importlib
# importlib.reload(mlreco.main_funcs)
1. Using YAML config file¶
First option is to define the configuration in a YAML text file or string, and load it directly into training or inference loop.
from mlreco.main_funcs import process_config, train, inference
import yaml
We will be using the LArCV data files in /gpfs/slac/staas/fs1/g/neutrino/kterao/data/dlprod_ppn_v10/combined directory, more specifically the files starting with train_512px. For more explanations on this configuration see section 2.
cfg = """
iotool:
batch_size: 4
shuffle: False
num_workers: 4
collate_fn: CollateSparse
sampler:
name: RandomSequenceSampler
batch_size: 4
dataset:
name: LArCVDataset
data_dirs:
- /gpfs/slac/staas/fs1/g/neutrino/kterao/data/dlprod_ppn_v10/combined
data_key: train_512px
limit_num_files: 10
schema:
input_data:
- parse_sparse3d_scn
- sparse3d_data
segment_label:
- parse_sparse3d_scn
- sparse3d_fivetypes
model:
name: uresnet
modules:
uresnet:
num_strides: 5
filters: 16
num_classes: 5
data_dim: 3
spatial_size: 512
network_input:
- input_data
loss_input:
- segment_label
training:
seed: -1
learning_rate: 0.001
gpus: '0'
weight_prefix: weights_trash/snapshot
iterations: 10
report_step: 1
checkpoint_step: 500
log_dir: log_trash
model_path: ''
train: True
debug: False
minibatch_size: -1
"""
cfg = yaml.load(cfg, Loader=yaml.Loader)
process_config(cfg)
if cfg['training']['train']:
train(cfg)
else:
inference(cfg)
2. Using API functions¶
from mlreco.trainval import trainval
from mlreco.main_funcs import make_directories
from mlreco.main_funcs import get_data_minibatched
from mlreco.main_funcs import process_config
from mlreco.iotools.factories import loader_factory
from mlreco.main_funcs import cycle
2.1. Training loop¶
Let's start by defining I/O configuration. We will be using the LArCV data files in /gpfs/slac/staas/fs1/g/neutrino/kterao/data/dlprod_ppn_v10/combined directory, more specifically the files starting with train_512px.
The schema field defines several data products. For each data product, the first string is the parser function name (e.g. parse_sparse3d_scn) and the following strings are TTree branch names to be fed to the parser.
The batch size needs to be defined in 2 places: batch_size field and in sampler.batch_size field.
io_cfg = {
"batch_size": 4,
"shuffle": False,
"num_workers": 4,
"collate_fn": "CollateSparse",
"sampler": {
"name": "RandomSequenceSampler",
"batch_size": 4
},
"dataset": {
"name": "LArCVDataset",
"data_dirs": ["/gpfs/slac/staas/fs1/g/neutrino/kterao/data/dlprod_ppn_v10/combined"],
"data_key": "train_512px",
"limit_num_files": 10,
"schema": {
"input_data": ["parse_sparse3d_scn", "sparse3d_data"],
"segment_label": ["parse_sparse3d_scn", "sparse3d_fivetypes"]
}
}
}
Then the network configuration, here we will train UResNet using SCN library, with depth 5 and 16 filters. We need to tell it the data spatial size and dimensions. The network_input and loss_input refer to I/O configuration names, they need to match between io_cfg and model_cfg.
model_cfg = {
"name": "uresnet",
"modules": {
"uresnet": {
"num_strides": 5,
"filters": 16,
"num_classes": 5,
"data_dim": 3,
"spatial_size": 512
}
},
"network_input": ["input_data"],
"loss_input": ["segment_label"]
}
Finally training configuration defines some global parameters such as learning rate and weights directory:
train_cfg = {
"seed": 123,
"learning_rate": 0.001,
"gpus": '0',
"weight_prefix": "weights_trash/snapshot",
"iterations": 10,
"report_step": 1,
"checkpoint_step": 500,
"log_dir": "log_trash",
"model_path": "",
"train": True,
"debug": False,
"minibatch_size": -1
}
The final configuration groups I/O, model and training:
cfg = {
"iotool": io_cfg,
"model": model_cfg,
"training": train_cfg
}
process_config(cfg)
Let's load the dataset using the configuration:
loader, cfg['data_keys'] = loader_factory(cfg)
dataset = iter(cycle(loader))
Initialize the trainer:
Trainer = trainval(cfg)
loaded_iteration = Trainer.initialize()
Create directories for weights and log files:
make_directories(cfg, loaded_iteration)
Now the actual training loop:
for i in range(10):
data_blob = get_data_minibatched(dataset, cfg)
res = Trainer.train_step(data_blob)
print(i, res)
Visualizing metrics¶
Now we want to visualize the loss and accuracy curves, so we load the CSV log file and plot these metrics.
log = np.genfromtxt('log_trash/train_log-0000000.csv', names=True, delimiter=',')
seaborn.set(rc={
'figure.figsize':(15, 10),
})
seaborn.set_context('talk') # or paper
plt.plot(log['iter'], log['loss_seg'])
plt.title("UResNet loss")
plt.xlabel("Iterations")
plt.ylabel("Segmentation Loss")
plt.plot(log['iter'], log['acc_seg'])
plt.title("UResNet nonzero accuracy")
plt.xlabel("Iterations")
plt.ylabel("Nonzero accuracy")
2.2. Inference¶
from mlreco.output_formatters import output
In order to record event displays in CSV files we need to specify output formatters:
model_cfg['outputs'] = ["input"]
The training configuration is almost same, we only need to change the train and model_path parameters. Let's use weights from a previous training:
inference_cfg = train_cfg.copy()
inference_cfg['gpus'] = '0'
inference_cfg["train"] = False
inference_cfg["model_path"] = "/gpfs/slac/staas/fs1/g/neutrino/ldomine/ppn_uresnet/weights_uresnet2/snapshot-999.ckpt"
We rebuild the configuration:
cfg = {
"iotool": io_cfg,
"model": model_cfg,
"training": inference_cfg
}
process_config(cfg)
Reload the dataset, reinitalize trainer:
loader, cfg['data_keys'] = loader_factory(cfg)
dataset = iter(cycle(loader))
Trainer = trainval(cfg)
loaded_iteration = Trainer.initialize()
make_directories(cfg, loaded_iteration)
Now run the inference loop for some iterations and use the output function to record the events in CSV files:
for i in range(10):
data_blob = get_data_minibatched(dataset, cfg)
res = Trainer.forward(data_blob)
print(i, res)
output(cfg['model']['outputs'], data_blob, res, cfg, i)
2.3. Event displays¶
Let's define some useful plotting functions using Plotly.
layout = go.Layout(margin=dict(l=0, r=0, b=0, t=0))
def get_trace(event, point_type, key="", point=False, cmax=0.05):
subevent = event[event['type'] == point_type]
if point:
trace1 = go.Scatter3d(
x=subevent['x'],
y=subevent['y'],
z=subevent['z'],
mode='markers',
marker=dict(
size=5,
color=subevent[key] if key else "orange",
#colorscale=[[0, "blue"], [0.2, "red"], [0.4, "firebrick"], [0.6, "lightseagreen"], [0.8, "mediumpurple"]],
colorscale="Rainbow",
opacity=0.8
)
)
else:
trace1 = go.Scatter3d(
x=subevent['x'],
y=subevent['y'],
z=subevent['z'],
mode='markers',
marker=dict(
size=1,
color=subevent["value"],
colorscale="Viridis",
opacity=0.8,
cmax=cmax,
cmin=0.
)
)
return trace1
def plot_event(data_products):
fig = go.Figure(data=data_products, layout=layout)
iplot(fig)
We load the output CSV file from one event:
event_id = 2
event = np.genfromtxt('log_trash/output-%.07d.csv' % event_id, names=True, delimiter=',')
Here is how it is recorded: we have 5 columns for each point: x, y, z for point coordinates, type to record different types of points (data, label, output, etc), and value (can be energy deposit, class label, etc).
event[:10]
For our purpose, all we need to know is energy deposits (type 0), segmentation labels (type 2) and later predictions (type 4)
2.3.1. Energy deposits¶
plot_event([get_trace(event, 0)])
2.3.2. Labels for UResNet¶
plot_event([get_trace(event, 2, cmax=5), get_trace(event, 1, key="value", point=True)])
2.4. Analysis and display of the network output¶
Now you might want to compute more metrics from the output of the network. To do this you need to specify a field analysis_keys in the model configuration. It will fetch the corresponding outputs from the network and feed it to our analysis function. Here we say that we want the first output of the network (index 0) to be saved under the key segmentation in the results dictionary:
model_cfg["analysis_keys"] = {
"segmentation": 0
}
Now we might also want to display the segmentation output of the network, so we will add the uresnet_ppn output formatter:
model_cfg['outputs'] = ["input", "uresnet_ppn"]
Now let's reload the configuration:
inference_cfg['gpus'] = '0'
cfg = {
"iotool": io_cfg,
"model": model_cfg,
"training": inference_cfg
}
process_config(cfg)
loader, cfg['data_keys'] = loader_factory(cfg)
dataset = iter(cycle(loader))
Trainer = trainval(cfg)
loaded_iteration = Trainer.initialize()
make_directories(cfg, loaded_iteration)
data_blob = get_data_minibatched(dataset, cfg)
res = Trainer.forward(data_blob)
print(res)
output(cfg['model']['outputs'], data_blob, res, cfg, 0)
Now you can see that in addition to loss and accuracy, res['segmentation'] contains the output of the semantic segmentation.
res['segmentation'][0].shape
You can now use it in your advanced analysis script, for example to compute accuracy per class, etc.
Let's see what the uresnet_ppn formatter recorded for us, the predictions:
event_id = 0
event = np.genfromtxt('log_trash/output-%07d.csv' % event_id, names=True, delimiter=',')
Predictions for this event¶
plot_event([get_trace(event, 4, cmax=5)]) # , get_trace(event1, 5, key="value", point=True)
Corresponding segmentation labels for this event¶
plot_event([get_trace(event, 2, cmax=5), get_trace(event, 1, key="value", point=True)])