onnx
====

.. py:module:: onnx

.. autoapi-nested-parse::

   onnx model management



Classes
-------

.. autoapisummary::

   onnx.OnnxRunner


Functions
---------

.. autoapisummary::

   onnx.get_prediction_function


Module Contents
---------------

.. py:class:: OnnxRunner(session, setup, device)

   .. py:attribute:: session


   .. py:attribute:: input_shape


   .. py:attribute:: output_shape


   .. py:attribute:: input_name


   .. py:attribute:: output_name


   .. py:attribute:: setup
      :type:  rex_xai.input.input_data.Setup


   .. py:attribute:: device
      :type:  str


   .. py:method:: run_on_cpu(tensors, target, raw, binary_threshold = None)

      Convert a pytorch tensor, or list of tensors, to numpy arrays on the cpu for onnx inference.



   .. py:method:: run_with_data_on_device(tensors, device, tsize, binary_threshold, raw=False, device_id=0, target=None)


   .. py:method:: gen_prediction_function()


.. py:function:: get_prediction_function(args)