diff --git a/TrkDiag/data/TrackPID_v1.onnx b/TrkDiag/data/TrackPID_v1.onnx new file mode 100644 index 0000000..51005e9 Binary files /dev/null and b/TrkDiag/data/TrackPID_v1.onnx differ diff --git a/TrkDiag/data/TrkQual_ANN1_v2.onnx b/TrkDiag/data/TrkQual_ANN1_v2.onnx new file mode 100644 index 0000000..7b89bcb Binary files /dev/null and b/TrkDiag/data/TrkQual_ANN1_v2.onnx differ diff --git a/TrkDiag/scripts/convert_trackpid_to_onnx.py b/TrkDiag/scripts/convert_trackpid_to_onnx.py new file mode 100644 index 0000000..7d221f1 --- /dev/null +++ b/TrkDiag/scripts/convert_trackpid_to_onnx.py @@ -0,0 +1,294 @@ +#!/usr/bin/env python3 +""" +Convert TrackPID_v1.dat format to ONNX model +Alternative version using TensorFlow/Keras which can export to ONNX + +The .dat file contains neural network weights in text format: + layer_name num_elements + value1 value2 ... valueN +""" + +import sys +import numpy as np +import argparse + + +def parse_dat_file(dat_filename): + """ + Parse the .dat file and extract layer weights and biases + Returns dict with layer names as keys and numpy arrays as values + """ + layers = {} + + with open(dat_filename, 'r') as f: + lines = f.readlines() + + i = 0 + while i < len(lines): + line = lines[i].strip() + + # Skip empty lines + if not line: + i += 1 + continue + + # Parse layer header: "layer_name num_elements" + parts = line.split() + if len(parts) >= 2: + layer_name = parts[0] + try: + num_elements = int(parts[1]) + except ValueError: + i += 1 + continue + + # Read values from next line(s) + values = [] + i += 1 + while len(values) < num_elements and i < len(lines): + value_line = lines[i].strip() + if value_line: + values.extend([float(x) for x in value_line.split()]) + i += 1 + + # Store as numpy array + if len(values) == num_elements: + layers[layer_name] = np.array(values[:num_elements], dtype=np.float32) + print(f"Loaded {layer_name}: shape {layers[layer_name].shape}") + else: + print(f"Warning: {layer_name} expected {num_elements} values but got {len(values)}") + else: + i += 1 + + return layers + + +def create_onnx_model_tf(layers, output_file): + """ + Create ONNX model using TensorFlow/Keras + + Network architecture (inferred from layer names): + Input (4) -> Dense(5) -> Dense(10) -> Dense(5) -> Dense(1) -> Output + """ + import tensorflow as tf + from tensorflow import keras + from tf2onnx.onnx_opset import onnxruntime_opset + import onnx + import tf2onnx.convert + + print("Using TensorFlow/Keras for model creation...") + + # Extract weights and biases for each layer + w1 = layers.get('tensor_sequential1dense1CastReadVariableOp0', np.zeros((4, 5), dtype=np.float32)) + b1 = layers.get('tensor_sequential1dense1BiasAddReadVariableOp0', np.zeros(5, dtype=np.float32)) + + w2 = layers.get('tensor_sequential1dense12CastReadVariableOp0', np.zeros((5, 10), dtype=np.float32)) + b2 = layers.get('tensor_sequential1dense12BiasAddReadVariableOp0', np.zeros(10, dtype=np.float32)) + + w3 = layers.get('tensor_sequential1dense21CastReadVariableOp0', np.zeros((10, 5), dtype=np.float32)) + b3 = layers.get('tensor_sequential1dense21BiasAddReadVariableOp0', np.zeros(5, dtype=np.float32)) + + w4 = layers.get('tensor_sequential1dense31CastReadVariableOp0', np.zeros((5, 1), dtype=np.float32)) + b4 = layers.get('tensor_sequential1dense31AddReadVariableOp0', np.array([0.0], dtype=np.float32)) + + # Reshape weights to proper dimensions + if w1.size == 20: + w1 = w1.reshape(4, 5) + if w2.size == 50: + w2 = w2.reshape(5, 10) + if w3.size == 50: + w3 = w3.reshape(10, 5) + if w4.size == 5: + w4 = w4.reshape(5, 1) + + print(f"W1 shape: {w1.shape}, B1 shape: {b1.shape}") + print(f"W2 shape: {w2.shape}, B2 shape: {b2.shape}") + print(f"W3 shape: {w3.shape}, B3 shape: {b3.shape}") + print(f"W4 shape: {w4.shape}, B4 shape: {b4.shape}") + + # Create Keras model + model = keras.Sequential([ + keras.layers.Dense(5, activation='relu', input_shape=(4,), weights=[w1, b1]), + keras.layers.Dense(10, activation='relu', weights=[w2, b2]), + keras.layers.Dense(5, activation='relu', weights=[w3, b3]), + keras.layers.Dense(1, weights=[w4, b4]) + ]) + + model.summary() + + # Convert to ONNX + print(f"Converting to ONNX and saving to {output_file}...") + spec = (tf.TensorSpec((None, 4), tf.float32, name="input"),) + output_path = output_file.replace('.onnx', '') + + model_proto, _ = tf2onnx.convert.from_keras(model, input_signature=spec, output_path=output_path) + + print(f"✓ Successfully converted to {output_file}") + + +def create_onnx_model_direct(layers, output_file): + """ + Create ONNX model directly without external dependencies + (using pure ONNX library if available) + """ + try: + import onnx + from onnx import helper, TensorProto + except ImportError: + print("Error: ONNX library not found. Install with: pip install onnx") + return False + + print("Using pure ONNX for model creation...") + + # Extract weights and biases + w1 = layers.get('tensor_sequential1dense1CastReadVariableOp0', np.zeros((4, 5), dtype=np.float32)) + b1 = layers.get('tensor_sequential1dense1BiasAddReadVariableOp0', np.zeros(5, dtype=np.float32)) + + w2 = layers.get('tensor_sequential1dense12CastReadVariableOp0', np.zeros((5, 10), dtype=np.float32)) + b2 = layers.get('tensor_sequential1dense12BiasAddReadVariableOp0', np.zeros(10, dtype=np.float32)) + + w3 = layers.get('tensor_sequential1dense21CastReadVariableOp0', np.zeros((10, 5), dtype=np.float32)) + b3 = layers.get('tensor_sequential1dense21BiasAddReadVariableOp0', np.zeros(5, dtype=np.float32)) + + w4 = layers.get('tensor_sequential1dense31CastReadVariableOp0', np.zeros((5, 1), dtype=np.float32)) + b4 = layers.get('tensor_sequential1dense31AddReadVariableOp0', np.array([0.0], dtype=np.float32)) + + # Reshape weights + if w1.size == 20: + w1 = w1.reshape(5, 4) + if w2.size == 50: + w2 = w2.reshape(10, 5) + if w3.size == 50: + w3 = w3.reshape(5, 10) + if w4.size == 5: + w4 = w4.reshape(1, 5) + + print(f"W1 shape: {w1.shape}, B1 shape: {b1.shape}") + print(f"W2 shape: {w2.shape}, B2 shape: {b2.shape}") + print(f"W3 shape: {w3.shape}, B3 shape: {b3.shape}") + print(f"W4 shape: {w4.shape}, B4 shape: {b4.shape}") + + # Create initializers + initializers = [] + + initializers.append(helper.make_tensor( + name="W1", + data_type=TensorProto.FLOAT, + dims=list(w1.shape), + vals=w1.flatten().tolist() + )) + initializers.append(helper.make_tensor( + name="B1", + data_type=TensorProto.FLOAT, + dims=list(b1.shape), + vals=b1.flatten().tolist() + )) + + initializers.append(helper.make_tensor( + name="W2", + data_type=TensorProto.FLOAT, + dims=list(w2.shape), + vals=w2.flatten().tolist() + )) + initializers.append(helper.make_tensor( + name="B2", + data_type=TensorProto.FLOAT, + dims=list(b2.shape), + vals=b2.flatten().tolist() + )) + + initializers.append(helper.make_tensor( + name="W3", + data_type=TensorProto.FLOAT, + dims=list(w3.shape), + vals=w3.flatten().tolist() + )) + initializers.append(helper.make_tensor( + name="B3", + data_type=TensorProto.FLOAT, + dims=list(b3.shape), + vals=b3.flatten().tolist() + )) + + initializers.append(helper.make_tensor( + name="W4", + data_type=TensorProto.FLOAT, + dims=list(w4.shape), + vals=w4.flatten().tolist() + )) + initializers.append(helper.make_tensor( + name="B4", + data_type=TensorProto.FLOAT, + dims=list(b4.shape), + vals=b4.flatten().tolist() + )) + + # Create input/output + input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 4]) + output_tensor = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 1]) + + # Create computation graph nodes + nodes = [ + helper.make_node("Gemm", inputs=["input", "W1", "B1"], outputs=["Y1"], alpha=1.0, beta=1.0, transB=1), + helper.make_node("Relu", inputs=["Y1"], outputs=["A1"]), + helper.make_node("Gemm", inputs=["A1", "W2", "B2"], outputs=["Y2"], alpha=1.0, beta=1.0, transB=1), + helper.make_node("Relu", inputs=["Y2"], outputs=["A2"]), + helper.make_node("Gemm", inputs=["A2", "W3", "B3"], outputs=["Y3"], alpha=1.0, beta=1.0, transB=1), + helper.make_node("Relu", inputs=["Y3"], outputs=["A3"]), + helper.make_node("Gemm", inputs=["A3", "W4", "B4"], outputs=["output"], alpha=1.0, beta=1.0, transB=1), + ] + + # Create graph and model + graph = helper.make_graph( + nodes=nodes, + name="TrackPID_Network", + inputs=[input_tensor], + outputs=[output_tensor], + initializer=initializers + ) + + model = helper.make_model(graph, producer_name="Mu2e_TrackPID_Converter", opset_imports=[helper.make_opsetid("", 12)]) + + # Save and verify + print(f"Saving to {output_file}...") + onnx.save(model, output_file) + + print("Verifying model...") + onnx.checker.check_model(model) + print("✓ Model verification passed!") + + return True + + +def main(): + parser = argparse.ArgumentParser(description="Convert TrackPID_v1.dat to ONNX format") + parser.add_argument("input", help="Input .dat file") + parser.add_argument("-o", "--output", help="Output .onnx file", default="TrackPID_v1.onnx") + parser.add_argument("--use-tf", action="store_true", help="Use TensorFlow/Keras for conversion (if available)") + + args = parser.parse_args() + + print(f"Reading {args.input}...") + layers = parse_dat_file(args.input) + + print(f"\nLoaded {len(layers)} layer(s)") + + if args.use_tf: + try: + create_onnx_model_tf(layers, args.output) + except ImportError as e: + print(f"TensorFlow not available ({e}), falling back to direct ONNX creation...") + create_onnx_model_direct(layers, args.output) + else: + success = create_onnx_model_direct(layers, args.output) + if not success: + print("\nTrying TensorFlow approach...") + try: + create_onnx_model_tf(layers, args.output) + except Exception as e: + print(f"TensorFlow approach failed: {e}") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/TrkDiag/src/SConscript b/TrkDiag/src/SConscript index 93181d8..efb348f 100644 --- a/TrkDiag/src/SConscript +++ b/TrkDiag/src/SConscript @@ -49,7 +49,8 @@ mainlib = helper.make_mainlib ( [ 'boost_filesystem', 'hep_concurrency', 'TMVA', - 'ROOTTMVASofie' + 'ROOTTMVASofie', + 'onnxruntime' ] ) # Fixme: split into link lists for each module. @@ -101,7 +102,8 @@ helper.make_plugins( [ 'hep_concurrency', 'TMVA', 'ROOTTMVASofie', -'pthread' + 'pthread', + 'onnxruntime' ] ) helper.make_dict_and_map( [ diff --git a/TrkDiag/src/TrackPID_module.cc b/TrkDiag/src/TrackPID_module.cc index 53d5c09..34f99f4 100644 --- a/TrkDiag/src/TrackPID_module.cc +++ b/TrkDiag/src/TrackPID_module.cc @@ -23,6 +23,8 @@ #include "ArtAnalysis/TrkDiag/inc/TrackPID.hxx" #include "Offline/GeometryService/inc/GeomHandle.hh" #include "Offline/CalorimeterGeom/inc/DiskCalorimeter.hh" +//ONNX +#include "onnxruntime/core/session/onnxruntime_cxx_api.h" // C++ #include #include @@ -65,6 +67,26 @@ namespace mu2e { int _debug; std::shared_ptr mva_; + + // Below is copied from TrackQuality_module.cc + Ort::Env _env; + Ort::SessionOptions _session_options; + Ort::Session _session; + Ort::AllocatorWithDefaultOptions _allocator; + Ort::AllocatedStringPtr _input_name; + Ort::TypeInfo _type_info; + Ort::ConstTensorTypeAndShapeInfo _tensor_info; + std::vector _input_shape; + size_t _total_size; + Ort::MemoryInfo _memory_info; + Ort::AllocatedStringPtr _output_name; + + std::string print_shape(const std::vector& v) { + std::stringstream ss(""); + for (std::size_t i = 0; i < v.size() - 1; i++) ss << v[i] << "x"; + ss << v[v.size() - 1]; + return ss.str(); + } }; TrackPID::TrackPID(const Parameters& conf) : @@ -73,12 +95,30 @@ namespace mu2e { _dtoffset(conf().DT()), _kalSeedPtrTag(conf().kalSeedPtrTag()), _printMVA(conf().printMVA()), - _debug(conf().debug()) + _debug(conf().debug()), + _env(ORT_LOGGING_LEVEL_WARNING, "ONNXInference"), + _session(_env, "ArtAnalysis/TrkDiag/data/TrkQual_ANN1_v2.onnx", _session_options), + _input_name(_session.GetInputNameAllocated(0, _allocator)), + _type_info(_session.GetInputTypeInfo(0)), + _tensor_info(_type_info.GetTensorTypeAndShapeInfo()), + _input_shape(_tensor_info.GetShape()), // Get input shape from model + _memory_info(Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault)), + _output_name(_session.GetOutputNameAllocated(0, _allocator)) { produces(); ConfigFileLookupPolicy configFile; mva_ = std::make_shared(configFile(conf().datFilename())); + // Handle dynamic dimensions if needed + for (auto& dim : _input_shape) { + if (dim == -1) { dim = 1; } // Set dynamic dims to 1 (or your desired value) + } + + // Calculate total size + _total_size = 1; + for (auto dim : _input_shape) { + _total_size *= dim; + } } void TrackPID::produce(art::Event& event ) { @@ -90,6 +130,8 @@ namespace mu2e { event.getByLabel(_kalSeedPtrTag, kalSeedPtrHandle); const auto& kalSeedPtrs = *kalSeedPtrHandle; + std::vector input_tensor_values(_total_size, 0.0f); // Initialize with zeros + // Go through the tracks and calculate the track PID for (const auto& kalSeedPtr : kalSeedPtrs) { const auto& kalSeed = *kalSeedPtr; @@ -114,14 +156,43 @@ namespace mu2e { // the following includes the (Calibrated) light-propagation time delay. It should eventually be put in the reconstruction FIXME! // This velocity should come from conditions FIXME! features[3] = tchs.t0().t0()-tchs.time()- std::min((float)200.0,std::max((float)0.0,tchs.hitLen()))*0.005 - _dtoffset; + + // For ONNX: + input_tensor_values[0] = features[0]; + input_tensor_values[1] = features[1]; + input_tensor_values[2] = features[2]; + input_tensor_values[3] = features[3]; + Ort::Value input_tensor = Ort::Value::CreateTensor(_memory_info, + input_tensor_values.data(), + input_tensor_values.size(), + _input_shape.data(), + _input_shape.size() + ); + // Run inference + const char* input_names[] = {_input_name.get()}; + const char* output_names[] = {_output_name.get()}; + auto output_tensors = _session.Run( + Ort::RunOptions{nullptr}, + input_names, + &input_tensor, + 1, + output_names, + 1 + ); + // Get output + float* output_data = output_tensors[0].GetTensorMutableData(); + // hard cut on the energy difference. This rejects cosmic rays which hit the calo and produce an upstream-going track that is then // reconstructed as a downstream particle associated to this cluster if(features[0] < _maxde){ // evaluate the MVA auto mvaout = mva_->infer(features.data()); mvaval = mvaout[0]; + } else { + output_data[0] = 0; // for ONNX, set the output to 0 if the energy difference is above the threshold } - else if (_debug > 0) { + + if (_debug > 0) { printf("energy difference at %f, above threshold %f", features[0], _maxde); } } diff --git a/TrkDiag/src/TrackQuality_module.cc b/TrkDiag/src/TrackQuality_module.cc index 2cc6468..e43a1ca 100644 --- a/TrkDiag/src/TrackQuality_module.cc +++ b/TrkDiag/src/TrackQuality_module.cc @@ -22,6 +22,8 @@ #include "Offline/RecoDataProducts/inc/KalSeed.hh" #include "Offline/RecoDataProducts/inc/MVAResult.hh" #include "ArtAnalysis/TrkDiag/inc/TrkQual_ANN1.hxx" +// ONNXRuntime +#include "onnxruntime/core/session/onnxruntime_cxx_api.h" // C++ #include #include @@ -66,20 +68,59 @@ namespace mu2e std::shared_ptr mva_; + Ort::Env _env; + Ort::SessionOptions _session_options; + Ort::Session _session; + Ort::AllocatorWithDefaultOptions _allocator; + Ort::AllocatedStringPtr _input_name; + Ort::TypeInfo _type_info; + Ort::ConstTensorTypeAndShapeInfo _tensor_info; + std::vector _input_shape; + size_t _total_size; + Ort::MemoryInfo _memory_info; + Ort::AllocatedStringPtr _output_name; + + std::string print_shape(const std::vector& v) { + std::stringstream ss(""); + for (std::size_t i = 0; i < v.size() - 1; i++) ss << v[i] << "x"; + ss << v[v.size() - 1]; + return ss.str(); + } }; TrackQuality::TrackQuality(const Parameters& conf) : art::EDProducer{conf}, _kalSeedPtrTag(conf().kalSeedPtrTag()), _printMVA(conf().printMVA()), - _debug(conf().debug()) + _debug(conf().debug()), + + _env(ORT_LOGGING_LEVEL_WARNING, "ONNXInference"), + _session(_env, "ArtAnalysis/TrkDiag/data/TrkQual_ANN1_v2.onnx", _session_options), + _input_name(_session.GetInputNameAllocated(0, _allocator)), + _type_info(_session.GetInputTypeInfo(0)), + _tensor_info(_type_info.GetTensorTypeAndShapeInfo()), + _input_shape(_tensor_info.GetShape()), // Get input shape from model + _memory_info(Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault)), + _output_name(_session.GetOutputNameAllocated(0, _allocator)) { produces(); ConfigFileLookupPolicy configFile; mva_ = std::make_shared(configFile(conf().datFilename())); - } + + // Handle dynamic dimensions if needed + for (auto& dim : _input_shape) { + if (dim == -1) { dim = 1; } // Set dynamic dims to 1 (or your desired value) + } + + // Calculate total size + _total_size = 1; + for (auto dim : _input_shape) { + _total_size *= dim; + } + } + void TrackQuality::produce(art::Event& event ) { // create output unique_ptr mvacol(new MVAResultCollection()); @@ -89,6 +130,9 @@ namespace mu2e event.getByLabel(_kalSeedPtrTag, kalSeedPtrHandle); const auto& kalSeedPtrs = *kalSeedPtrHandle; + // Prepare input tensor + std::vector input_tensor_values(_total_size, 0.0f); // Initialize with zeros + // Go through the tracks and calculate their track qualities for (const auto& kalSeedPtr : kalSeedPtrs) { const auto& kalSeed = *kalSeedPtr; @@ -135,6 +179,11 @@ namespace mu2e features[3] = (double) nnullambig / nactive; features[4] = kalSeed.fitConsistency(); features[6] = (double)nmatactive / nactive; + input_tensor_values[0] = nactive; + input_tensor_values[1] = (double) nactive / nhits; + input_tensor_values[3] = (double) nnullambig / nactive; + input_tensor_values[4] = kalSeed.fitConsistency(); + input_tensor_values[6] = (double)nmatactive / nactive; // Now get the features that are for the entrance of the trackre bool entrance_found = false; @@ -143,6 +192,8 @@ namespace mu2e if (kinter.surfaceId() == SurfaceIdDetail::TT_Front) { // we only want the tracker entrance (sid=0) features[2] = sqrt(kinter.loopHelix().paramVar(KinKal::LoopHelix::t0_)); features[5] = kinter.momerr(); + input_tensor_values[2] = sqrt(kinter.loopHelix().paramVar(KinKal::LoopHelix::t0_)); + input_tensor_values[5] = kinter.momerr(); entrance_found = true; break; } @@ -150,18 +201,41 @@ namespace mu2e if (!entrance_found) { features[2] = -9999; features[5] = -9999; + input_tensor_values[2] = -9999; + input_tensor_values[5] = -9999; } std::vector mvaout = mva_->infer(features.data()); + Ort::Value input_tensor = Ort::Value::CreateTensor(_memory_info, + input_tensor_values.data(), + input_tensor_values.size(), + _input_shape.data(), + _input_shape.size() + ); + // Run inference + const char* input_names[] = {_input_name.get()}; + const char* output_names[] = {_output_name.get()}; + auto output_tensors = _session.Run( + Ort::RunOptions{nullptr}, + input_names, + &input_tensor, + 1, + output_names, + 1 + ); + // Get output + float* output_data = output_tensors[0].GetTensorMutableData(); + if (!entrance_found) { mvaout[0] = 0; // this is not a good track + output_data[0] = 0; } if(_debug > 0) { - printf("[TrackQuality::%s::%s] Inputs = %.0f, %.4f, %.4f, %.4f, %.4f, %.4f %.4f --> output = %.4f\n", + printf("[TrackQuality::%s::%s] Inputs = %.0f, %.4f, %.4f, %.4f, %.4f, %.4f %.4f --> output = %.4f (ORT: %.4f)\n", __func__, moduleDescription().moduleLabel().c_str(), - features[0], features[1], features[2], features[3], features[4], features[5], features[6], mvaout[0]); + features[0], features[1], features[2], features[3], features[4], features[5], features[6], mvaout[0], output_data[0]); } mvacol->push_back(MVAResult(mvaout[0]));