peal.architectures.onnx_predictor.select_onnx_outputs

peal.architectures.onnx_predictor.select_onnx_outputs(src, dst, indices, output_name='output')[source]

Write a copy of the classifier at src whose single output is logits[:, indices] (a Gather on the last axis), so an N-way model becomes the binary model PEAL analyses. Returns dst.

The corrected model DiDAE exports afterwards is binary as well: it is the finetuned two-class head the user asked to analyse, not the original N-way graph.

Parameters:
  • src (str) – Source .onnx file.

  • dst (str) – Destination .onnx file (written and checked with onnx.checker).

  • indices (sequence of int) – At least two distinct output indices, in the order the new outputs should have.

  • output_name (str) – Name of the new graph output (suffixed _selected when the old output already has that name).

Returns:

dst.

Return type:

str

Raises:

ValueError – If fewer than two distinct indices are given or an index is out of range for a statically shaped output.