Noob Question: Can we export a Compute Graph to ONNX with Candle?

[nihalpasham]

I’m not well-versed in the different ML storage/ graph formats, so this may be a noob question.

• Does candle provide a way to save (or export) a (very simple) compute graph as an ONNX file? To be more precise, I have a topologically ordered vector of tensors (constructed as shown below). I saved it as a .safetensors file (which does not contain any graph level info), but I was wondering if there is a way to convert it to the ONNX format.

let bsize = 1usize;
let x = 8usize;
let y = 512usize;
let device = Device::new_metal(0).unwrap();
let vmap = VarMap::new();
let vb = VarBuilder::from_varmap(&vmap, DType::F32, &device);
let ll = linear(512, 512, vb.pp("ll")).unwrap();

let random = Tensor::randn(0f32, 1., (bsize, x, y), &device).unwrap();
let res = ll.forward(&random).unwrap();
let ordered_nodes = res.sort_nodes();
Tensor::get_op_graph(ordered_nodes);

[nihalpasham]

Tagging @shua, @B1rtek, @mokulus, and @LaurentMazare, i.e., contributors to the candle-onnx crate, for any inputs. I would appreciate it if anyone could point me in the right direction here.

[shua]

For your specific case of a list of tensors, I imagine you could build a graph out of "Constant" node types, and then write that out to a file. Something like

let mut g = GraphProto::default();
for (i, t) in tensors.iter().enumerate() {
  g.node.push(NodeProto{
    op_type: "Constant".to_string(),
    attribute: vec![AttributeProto{
      name: "value".to_string(),
      r#type: AttributeType::Tensor as i32,
      t: ... map data from candle::Tensor to onnx::TensorProto ...,
      ..default()
    }],
    output: vec![i.to_string()],
    ..default()
  });
  g.output.push(ValueInfoProto{name: i.to_string(), ..default() })
}
let mut buf = vec![];
let model = ModelProto{ graph: Some(g), ..default() }
model.encode(&mut buf).unwrap();
let f = std::fs::File::create("path/to/model.onnx");
f.write_all(&buf).unwrap()

This could work for a list of tensors, not really if you build anything more complex.

[nihalpasham]

Thank you for taking the time to look into this. In my case, I put together a simple unit test to print the list of tensors (which is actually an op_graph for the above example) by reusing the existing topological ordering implementation baked into Candle.

The results are as follows:

cargo test test_get_op_graph -- --nocapture
   Compiling optimus v0.1.0 (/Users/nihal.pasham/devspace/metal/optimus)
    Finished test [unoptimized + debuginfo] target(s) in 0.50s
     Running unittests src/main.rs (target/debug/deps/optimus-fbefd677e69a344d)

running 1 test
BinaryOp: Add, arg1: TensorId(6), arg2: TensorId(7)
BroadcastOp: arg: TensorId(2)
None: arg: TensorId(2)
MatmulOp: arg1: TensorId(3), arg2: TensorId(5)
TransposeOp: (1, 2) dims, arg: TensorId(4)
BroadcastOp: arg: TensorId(1)
None: arg: TensorId(1)

I presume that with just constant nodes, we wouldn't be able to perform any actual computation as the corresponding ONNX file would be a static graph of inputs and outputs. However, your above comment leads me to believe that we can extend our ONNX graph construction to include more trivial operators (such as add, matmul, transpose, etc.). Is my understanding correct?

[shua]

yeah sure, you can build whatever graph you want from the onnx proto types (eg NodeProto, etc), and then encode that model to a file as sketched out above. It's pretty verbose, but it should be possible.

[shua]

I am also a noob here, but my guess is that onnx export is not currently supported. This is based on looking at the public api for candle_onnx (v0.5.1) it only contains a read_file for loading onnx files, but not some similar write_file for saving them.