How does tract_onnx::tract_hir::ops::cnn::Conv implement Op ?

[alexander-camuto]

Hello !

Have been loving using the library so far. I'm currently loading up a simple single layer convolutional model, saved using pytorch. I then load up the saved model as such:

let model = tract_onnx::onnx().model_for_path(path).unwrap();

I then iterate over the model nodes, call op() on each one and print the op name. As one might expect I get the following output:

Source
Const
Const
ConvHir

The first Source corresponds to the model input, the two Const to the layer kernel and bias, and finally ConvHir to the conv specific parameters. I specifically want to extract information about the padding and stride used in the final ConvHir node. I've tried to do this as such:

let op = node.op();

As expected if I print op I get the following output,

Conv { data_format: NCHW, kernel_fmt: OIHW, dilations: Some([1, 1]), kernel_shape: Some([5, 5]), padding: Explicit([1, 1], [1, 1], false), strides: Some([1, 1]), group: Some(1), x_scale_input: None, x_zero_point_input: None, k_input: None, k_scale_input: None, k_zero_point_input: None, y_scale_input: None, y_zero_point_input: None, bias_input: Some(2), override_output_datum_type: None }

However the op() method returns &dyn Op . So I can't directly access the padding and stride methods of the Conv op without explicitly casting it to tract_onnx::tract_hir::ops::cnn::Conv, which I initially assumed was the correct type here.
I try and do so with the provided methods on the Op trait,

let conv_node: &Conv = match op.downcast_ref() {
                    Some(b) => b,
                    None => panic!("op isn't a Conv!"),
                };

but get the following error: the trait "Op" is not implemented for "Conv".
It seems like tract_onnx::tract_hir::ops::cnn::Conv is the wrong type to try and downcast to, but looking through the code I can't find any other struct that would be suitable here.

As a sanity check I also run

assert!(op.is::<Conv>());

which also panics and fails.

Would appreciate a pointer to which type I should be downcasting to.

[kali]

Conv is defined in hir/src/ops/cnn/conv.rs. It does not directly implements Op, it is an Expansion instead. Defined in hir/src/ops/expandable.rs. It is wrapped in Box<dyn Expansion> which does implements Op and EvalOp in a very inneficient way (by generic an equivalent network to eval itself).

When you call into_typed (directly or through into_declutter or into_optimized) the expansion expands to actualy convolution operators (most likely ConvUnary from tract-core).

[alexander-camuto]

Awesome that did it. If anyone comes up against this:

let op = Box::new(node.op());

let conv_node: &Conv = match op.downcast_ref::<Box<dyn Expansion>>() {
     Some(b) => match (*b).as_any().downcast_ref() {
             Some(b) => b,
              None => panic!("not a conv!"),
    },
     None => panic!("op isn't an Expansion!"),
 };