peal.architectures.onnx_predictor.OnnxPredictor

class peal.architectures.onnx_predictor.OnnxPredictor[source]

Bases: Module

A trainable torch module converted from an ONNX classifier.

forward takes the same [B, C, H, W] float tensor the ONNX graph declared and returns its first output. source and input_shape are kept so the corrected model can be exported with the interface the user uploaded.

Parameters:

  • graph_module (torch.fx.GraphModule) – The module onnx2torch.convert produced.

  • source (str, optional) – Absolute path of the ONNX file the module came from.

  • input_name (str) – Name of the graph input, kept for re-export.

  • input_shape (list of int or None, optional) – Declared input shape (symbolic dims as None); used by export_to_onnx() when no shape is given.

Methods

forward

Run the converted graph and return its first output tensor (in chunks when the graph declares a fixed batch size).
__init__(graph_module, source=None, input_name='input', input_shape=None)[source]

Wrap graph_module and remember the ONNX interface it came from.

forward(x)[source]

Run the converted graph and return its first output tensor (in chunks when the graph declares a fixed batch size).