PyTorch.Geometric to HLS4ML

Created by: abdelabd

Purpose

Add pyg_to_hls converter to parse an Interaction Network PyTorch Geometric model which is composed of alternating EdgeBlocks (which applies a neural network to "edge features" to compute messages), Aggregate layers (which aggregates edge features for nodes that are connected to it), and NodeBlocks (which applies a neural network aggregated edge features.

PyTorch Geometric models expect input data consisting of node attributes node_attr of size [n_nodes, n_node_features], edge attributes edge_attr of size [n_edges, n_edge_features], and an edge index edge_index of size [2, n_edges], which specifies the data's connectivity.

The currently supported network for charged particle tracking output edge weights of size [n_edges, 1], which classify the edges as valid (part of a good track segment) or not.

Slides

https://docs.google.com/presentation/d/1C3EuVq4FaaayNSBeaRBHyHnWcujLbXASxAR2cqQyLVo/edit?usp=sharing

Test script

Notebook to test the conversion: https://github.com/abdelabd/manual_GNN_conversion/blob/main/pyg_to_hls_walkthrough.ipynb Script to test the conversion: https://github.com/abdelabd/manual_GNN_conversion/blob/main/convert_model.py

python convert_model.py conversion_config.yaml --n-graphs=100 --aggregation add --flow source_to_target --precision 'ap_fixed<16,8> --max-nodes=28 --max-edges=37 --n-neurons=8

or more fully featured code: https://github.com/abdelabd/manual_GNN_conversion

python test_model.py test_config.yaml --n-graphs 100 --aggregation add --flow source_to_target --precision 'ap_fixed<16,8>' --max-nodes 28 --max-edges 37 --n-neurons 8 --synth

Changes

Frontend

• added onto hls4ml.converters.__init__.py:
  • added from hls4ml.converters.pyg_to_hls import pyg_to_hls
  • added convert_from_pyg_model() function
  • registered the pyg block_handlers
• added hls4ml/converters/pyg/__init__.py
  • nothing here
• added hls4ml/converters/pyg/interaction_network_blocks.py
  • has the pyg block_handlers (analogous to layer_handlers) for NodeBlock, EdgeBlock, Aggregate
• added hls4ml.converters.pyg_to_hls.py
  • has the class PygModelReader(PyTorchModelReader)
  • has the function pyg_to_hls()
  • has the pyg 'block_handlers'
• added onto hls4ml.utils.config.py
  • added config_from_pyg_model() function
• changed hls4ml.writer.vivado_writer.py
  • changed VivadoWriter.write_defines() method
    • if certain macros are shared between different layers of an HLSModel, then these macros are not written multiple times in the defines.h file.
  • changed VivadoWriter.write_yml() method
    • If yaml.dump() has trouble writing a torch model with submodules composed of torch models themselves, then it writes the location of the model's state dictionary instead.
• added onto hls4ml.model.hls_layers
  • added class GraphBlock(Layer): parent class for EdgeBlock and NodeBlock. Allows hls4ml to handle an EdgeBlock or NodeBlock as if it were a single layer (like Dense, Convolution, etc.), rather than a submodule composed of several layers.
  • added class EdgeBlock(GraphBlock)
  • added class NodeBlock(GraphBlock)
  • added class Aggregate(Layer)
• hls4ml.model.hls_model
  • changed HLSModel.get_weights_data()
    • handling for getting weights data if the layer in question doesn't belong directly to the torch model in question, but rather to a submodule of the torch model such as an EdgeBlock or NodeBlock

Backend

• added onto nnet_utils/nnet_array.h
  • added struct matrix_config which specifies parameters for conversion functions
  • added vec_to_mat() function to convert from 1D unrolled arrays and 2D matrices
  • added mat_to_vec() function to convert from 2D matrices to 1D unrolled arrays
• changed nnet_utils/nnet_dense.h and changed nnet_utils/nnet_dense_resource.h
  • added parameter "static const bool remove_pipeline_pragma = false;" to the dense_config
  • under the ReuseLoop for the dense_resource_rf_leq_nin(), dense_resource_rf_gt_nin_rem0(), and dense_resource_rf_gt_nin() functions: "#pragma HLS PIPELINE II=1 rewind" --> if remove_pipeline_pragma=false then #pragma HLS PIPELINE II=1 rewind.
    • This change allows the pipelining to happen at the GraphBlock layer level instead of the individual dense layer level
• added nnet_utils/nnet_graph.h
  • has the backend code for the EdgeBlock, NodeBlock, and Aggregate layers.