mortont/axon_onnx

`ReduceSum` not implemented

Closed this issue · 2 comments

adri commented

Hi! Thanks a lot for this library. I'm new to deep learning and wanted to try out a few models in Elixir. Specifically I'm interested in trying out semantic search by using embeddings from sentence transformers (sbert) using pre-trained models.

I managed to export a model all-MiniLM-L6-v2 to onnx.

I run the following code to load the model:

{model, params} = AxonOnnx.import("priv/models/sentence-transformers-all-MiniLM-L6-v2.onnx", batch: 1, sequence: 1)

Using {:axon_onnx, github: "elixir-nx/axon_onnx"}

** (ArgumentError) unsupported "ReduceSum"
    (axon_onnx 0.3.0) lib/axon_onnx/deserialize.ex:2189: AxonOnnx.Deserialize.recur_nodes/2
    (elixir 1.14.1) lib/enum.ex:2468: Enum."-reduce/3-lists^foldl/2-0-"/3
    (axon_onnx 0.3.0) lib/axon_onnx/deserialize.ex:43: AxonOnnx.Deserialize.graph_to_axon/2
    (axon_onnx 0.3.0) lib/axon_onnx/deserialize.ex:26: AxonOnnx.Deserialize.to_axon/2
    iex:1: (file)

I guess I'd need to implement the ReduceSum method. Happy to help, but I have no idea where to start?

adri commented

I found the ReduceSum specification, maybe it's not too hard. Sorry for my beginner questions.

  • Could I just use Nx.sum for that?
  • How do I test if the implementation works? I see some commented out tests but nothing is about ReduceSum

You can just wrap an Nx.sum in an Axon layer, but you'll need to ensure you get axes using the Axon constant! helper. You won't be able to run the test suite as ONNX passes those as runtime values :(