Converting a model with a matrix multiplication (PyTorch)

[ibengtsson]

Hi! I have a fairly simple network written in PyTorch but it requires several matrix multiplications in the forward pass. As this is not a functionality currently supported, I opted for writing a custom layer simply wrapping a torch.matmul operation using the Extension API. However, this does not seem to work... When trying to convert the model, I still get: "Unsupported function matmul".

Might it be that I lack understanding and it is not possible to do convert a model with a matrix multiplication in the way I intended? Or could the problem be caused by how the symbolic trace works when checking what PyTorch modules the backend supports? The example on the Extension API uses Keras, perhaps using PyTorch introduces problems I was not aware of?

I'm very thankful for any pointers on how to move forward with the issue!

[vloncar]

Well, that is a bit unrelated as the question is neither about TorchScript nor ONNX.

@itzakkk The extension API is not yet supported well with PyTorch, we're still working on it. However, you can still achieve what you want by creating a custom function and decorating it with torch.fx.wrap (see the docs). You can then follow the rest of the extension API flow and register a handler with hls4ml.converters.pytorch_to_hls.register_pytorch_layer_handler (it's not yet public API). Note that we don't have optimized matrix multiplication, we use matrix-vector multiplication, so you'll need to provide HLS for that yourself.

[ibengtsson]

Thank you for the kind answer @EdmTechPhoenix , but as @vloncar says I'm confident my issue isn't related to TorchScript or ONNX.

Your answer was however very helpful @vloncar ! I really appreciate it. Following your suggestions I managed to get a small toy model to work without too much trouble.

[vloncar]

Care to share your toy example? 😄 We're always looking of examples of what people try to do with extensions, this helps us make hls4ml more extensible.

[ibengtsson]

Certainly! My network is actually a graph neural network, but only uses a simple graph convolution operation given by
$$x_i' = W_1 x_i + W_2 \sum_{j \in \mathcal{N}(i)} \epsilon_{ij} \cdot x_j.$$
In the original model the graph convolution was implemented in PyTorch Geometric, but by replacing the summation with a matrix multiplication of the adjacency matrix (filled with edge weights instead of binary labels) and the node features I figured it was possible to implement it as a linear combination of torch.nn.Linear. More concretely, I'm implementing

class CustomGraphConv(nn.Module):
    def __init__(self, input_features, output_features):
        super().__init__()
        self.weight_node = nn.Linear(input_features, output_features, bias=False)
        self.weight_adj = nn.Linear(input_features, output_features, bias=True)

    def forward(self, x, adj):
        node_term = self.weight_node(x)
        adjaceny_sum = pmat_mul(adj, x)
        adjaceny_term = self.weight_adj(adjaceny_sum)

        return node_term + adjaceny_term

and need an implementation of pmat_mul that hls4ml can parse. I have not yet had time to verify the results, but by decorating a custom function with torch.fx.wrap I was able to run hls4ml.converters.convert_from_pytorch_model at least!