peal.architectures.onnx_predictor

ONNX classifiers as PEAL predictors.

PEAL accepted .onnx predictors before, but only as an onnxruntime closure: fine for the analysis half (DiDAE’s sweep, CFKD’s counterfactual search only ever query the student), useless for the correction half, whose finetuning needs an nn.Module with parameters. This module closes that gap:

  • load_onnx_as_torch() converts the graph with onnx2torch into a trainable OnnxPredictor (an nn.Module that pickles, deep-copies and re-exports), after normalising the graph so the converter accepts what torch.onnx.export produces (weights aliased through Identity nodes, Constant nodes instead of initializers). The converted module is checked against onnxruntime on a random batch before it is trusted.

  • load_onnx_runtime() is the old closure, kept as the fallback for graphs onnx2torch cannot convert, with a CPU provider fallback instead of the hard-coded CUDA provider (the CPU-only wheel is what pip gives you on aarch64 hosts such as the DGX Spark).

  • export_to_onnx() writes a finetuned predictor back out, so a user who uploaded an ONNX file gets a corrected ONNX file back.

  • inspect_onnx() reads the declared input/output shapes, which the web demo uses to validate a user’s mandatory fields against the graph.

Set PEAL_ONNX_CONVERT=0 to skip the conversion and always use the runtime closure (the behaviour every pre-2026-09-25 config saw).

Functions

export_to_onnx

Export a predictor to ONNX with a dynamic batch axis.

find_final_linear

Locate the classifier's final linear layer in an ONNX graph.

inspect_onnx

Declared interface of an ONNX graph.

load_onnx_as_torch

Convert an ONNX classifier into an OnnxPredictor.

load_onnx_predictor

What get_predictor calls for a .onnx path: the trainable module when the conversion succeeds, the runtime closure otherwise.

load_onnx_runtime

The pre-conversion behaviour: an onnxruntime closure (not trainable).

normalize_onnx_graph

Rewrite an ONNX model in place so onnx2torch accepts it.

select_onnx_outputs

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.

truncate_to_features

Write a copy of src whose only output is feature_tensor (the input of the final linear layer found by find_final_linear()).

write_final_linear

Replace the final layer's weights for the selected classes.

Classes

OnnxPredictor

A trainable torch module converted from an ONNX classifier.