Skip to content

MF Onnx and cleaner output predictions

Nikita Ivvan Pond requested to merge npond/salt:mf_onnx into main

Description

Getting MF to work for onnx export:

  • Modifies the main MF code slightly to allow for times when we have 0 tracks
  • Fix some of the lines that were causing other issues for onnx
  • Implements onnx outputs, currently we write the 'leading' object as scalers to decorate the jet. Here leading is defined as the object with the highest predicted regression[0], which is pT for the vertexing case
  • Modified the mask_to_indices function to return indices in the same format as GN2 vertexing. That is, first find any tracks which the model defines as a vertex. Then, for all remaining non-padded tracks, use a consecutively increasing index. E.g 1 mask [False, True, True, False, False] -> [1, 0, 0, 2, 3] as we have vertex 0 at idx=1,2 and all remaining are not part of a vertex.
  • Also modify the prediction writer to output in the same format.

This MR relates to #53. New issue added - need to be eventually able to write all regression values, rather than just leading regression. But this will require some work.

Review checklist:

  • CI Passing
  • Comments addressed
  • Source branch is up to date with target
Edited by Nikita Ivvan Pond

Merge request reports