Dummy PR to test #379

hls4ml/converters/__init__.py

@@ -11,6 +11,7 @@
 #----------Make converters available if the libraries can be imported----------#       
 try:
     from hls4ml.converters.pytorch_to_hls import pytorch_to_hls, get_supported_pytorch_layers, register_pytorch_layer_handler
+    from hls4ml.converters.pyg_to_hls import pyg_to_hls, get_supported_pyg_blocks, register_pyg_block_handler
     __pytorch_enabled__ = True
 except ImportError:
     warnings.warn("WARNING: Pytorch converter is not enabled!")
@@ -31,7 +32,7 @@
     __tensorflow_enabled__ = False
 
 #----------Layer handling register----------#
-model_types = ['keras', 'pytorch', 'onnx']
+model_types = ['keras', 'pytorch', 'onnx', 'pyg']
 
 for model_type in model_types:
     for module in os.listdir(os.path.dirname(__file__) + '/{}'.format(model_type)):
@@ -52,6 +53,8 @@
                             register_pytorch_layer_handler(layer, func)
                         elif model_type == 'onnx':
                             register_onnx_layer_handler(layer, func)
+                        elif model_type == 'pyg':
+                            register_pyg_block_handler(layer, func)
 
         except ImportError:
             continue
@@ -267,6 +270,134 @@ def convert_from_pytorch_model(model, input_shape, output_dir='my-hls-test', pro
 
     return pytorch_to_hls(config)
 
+def check_forward_dict(model, forward_dictionary):
+    for key in forward_dictionary:
+        try:
+            block = getattr(model, key)
+        except AttributeError:
+            raise AttributeError(f'Model is missing module "{key}" that is present in the provided forward dictionary; Check compatability')
+
+def convert_from_pyg_model(model, forward_dictionary, n_node, node_dim,
+                           n_edge, edge_dim, activate_final=None,
+                           output_dir='my-hls-test', project_name='myproject',
+                           fpga_part='xcku115-flvb2104-2-i', clock_period=5, io_type='io_parallel', hls_config={}):
+    check_forward_dict(model, forward_dictionary)
+    """
+
+    Convert a Pytorch.Geometric model to an hls model.
+
+    Parameters
+    ----------
+    model : Pytorch.geometric model object.
+        Model to be converted to hls model object.
+    n_node, n_edge: int, int
+        These parameters define the size of the graphs that your hls GNN 
+        accepts as input. Inputs must be truncated or zero-padded to this 
+        size before feeding them to your model. This is necessary because 
+        each layer of the hls/hardware implementation has a fixed size 
+        and cannot be resized. 
+    node_dim, edge_dim: int, int
+        node_dim defines the length of the vector used to represent each 
+        node in the graph-input. For example, if each node is represented 
+        as a 1x3 vector, node_dim=3. 
+        Likewise, edge_dim defines the length of the vector used to 
+        represent each edge in the graph-input.
+
+    forward_dictionary: OrderedDict object of the form {string: string}
+        Use this dictionary to define the order in which your model's
+        forward() method calls on the model's submodules. The keys
+        of the dictionary should be the names of your model's submodules, and the 
+        value stored in each key should indicate whether that submodule is an 
+        'EdgeBlock' (i.e. it predicts messages/edge-updates) or whether its a
+        'NodeBlock' (i.e. it predicts node-updates). 
+
+        For example, consider this InteractionNetwork (https://github.com/GageDeZoort/interaction_network_paper/blob/pytorch_geometric/models/interaction_network.py),
+        whose forward() method calls on its submodules in the following order:
+        1. An EdgeBlock named 'R1'
+        2. A NodeBlock named 'O'
+        3. An EdgeBlock named 'R2'
+
+        One would define its forward dictionary as such:
+        >>> forward_dictionary = OrderedDict()
+        >>> forward_dictionary['R1'] = 'EdgeBlock'
+        >>> forward_dictionary['O'] = 'NodeBlock'
+        >>> forward_dictionary['R2'] = 'EdgeBlock'
+
+        It is really important to define the submodules in the same order with which the 
+        forward() method calls on them. hls4ml has no other way of inferring this order. 
+
+    activate_final: string, optional 
+        If the activation of the final output is not already a layer in the corresponding
+        submodule, name the type of the activation function here. In the preceding example, 
+        one would pass the value 'sigmoid', because the final output of the model 
+        is the sigmoid-activated output of 'R2' (the last submodule called by the
+        forward() method). In other words, the model returns torch.sigmoid(self.R2(m2)). 
+        Other accepted values for this parameter include: 
+                                ['linear', 'relu', 'elu', 'selu', 'prelu', 'leaky_relu', 'softmax', 'tanh', 'softplus',  
+                                'softsign', 'hard_sigmoid','thresholded_relu', 'binary_tanh', 'ternary_tanh']
+    output_dir : string, optional
+        Output directory to write hls codes.
+    project_name : string, optional
+        hls project name.
+    fpga_part : string, optional
+        The particular FPGA part number that you are considering.
+    clock_period : int, optional
+        The clock period, in ns, at which your algorithm runs.
+    io_type : string, optional
+        Your options are 'io_parallel' or 'io_serial' where this really 
+        defines if you are pipelining your algorithm or not.
+    hls_config : dict, optional
+        Additional configuration dictionary for hls model.
+
+    Returns
+    -------
+    hls_model : hls4ml model object.
+
+    See Also
+    --------
+    hls4ml.convert_from_pytorch_model, hls4ml.convert_from_keras_model, 
+    hls4ml.convert_from_onnx_model
+
+    Example
+    --------
+    >>> import hls4ml
+    >>> config = hls4ml.utils.config_from_pyg_model(model, granularity='model')
+    >>>
+    >>> forward_dictionary = OrderedDict()
+    >>> forward_dictionary['R1'] = 'EdgeBlock'
+    >>> forward_dictionary['O'] = 'NodeBlock'
+    >>> forward_dictionary['R2'] = 'EdgeBlock'
+    >>> graph_dimensions = {"n_node": 112, "node_dim": 3, "n_edge": 148, "edge_dim": 4}
+    >>> hls_model = hls4ml.converters.convert_from_pyg_model(model, forward_dictionary,
+                                                             **graph_dimensions,
+                                                             activate_final='sigmoid'
+                                                             hls_config=config)
+
+    """
+
+    config = create_vivado_config(
+        output_dir=output_dir,
+        project_name=project_name,
+        fpga_part=fpga_part,
+        clock_period=clock_period,
+        io_type=io_type
+    )
+
+    config['PytorchModel'] = model
+    config['InputShape'] = {
+        'NodeAttr': [n_node, node_dim],
+        'EdgeAttr': [n_edge, edge_dim],
+        'EdgeIndex': [n_edge, 2]
+    }
+    config['ForwardDictionary'] = forward_dictionary
+    config['ActivateFinal'] = activate_final
+
+    model_config = hls_config.get('Model', None)
+    config['HLSConfig']['Model'] = _check_model_config(model_config)
+
+    _check_hls_config(config, hls_config)
+
+    return pyg_to_hls(config)
 
 def convert_from_onnx_model(model, output_dir='my-hls-test', project_name='myproject',
     fpga_part='xcku115-flvb2104-2-i', clock_period=5, io_type='io_parallel', hls_config={}):

hls4ml/converters/pyg/interaction_network_blocks.py

@@ -0,0 +1,60 @@
+import numpy as np
+from hls4ml.converters.pyg_to_hls import pyg_handler
+
+def parse_GraphBlock(block_name, config, n_node, n_edge, node_dim, edge_dim):
+    layer_dict = {
+        "name": block_name,
+        "n_node": n_node,
+        "n_edge": n_edge,
+        "node_dim": node_dim,
+        "edge_dim": edge_dim,
+    }
+
+    # get n_layers, out_dim
+    model = config['PytorchModel']
+    torch_block = getattr(model, block_name)
+    try:
+        torch_layers = torch_block.layers._modules
+    except AttributeError:
+        torch_layers = torch_block._modules
+
+    lcount = 0
+    for lname, l in torch_layers.items():
+        if l.__class__.__name__=="Linear":
+            lcount += 1
+            last_layer = l
+    layer_dict["n_layers"] = lcount
+    layer_dict["out_dim"] = last_layer.out_features
+    return layer_dict
+
+@pyg_handler('NodeBlock')
+def parse_NodeBlock(block_name, config, update_dict, index, n_node, n_edge, node_dim, edge_dim):
+    layer_dict = parse_GraphBlock(block_name, config, n_node, n_edge, node_dim, edge_dim)
+    layer_dict["class_name"] = "NodeBlock"
+    layer_dict["inputs"] = [update_dict["last_node_update"], update_dict["last_edge_aggr_update"]]
+    layer_dict["outputs"] = [f"layer{index}_out"]
+    update_dict["last_node_update"] = f"layer{index}_out"
+    return layer_dict, update_dict
+
+@pyg_handler('EdgeBlock')
+def parse_EdgeBlock(block_name, config, update_dict, index, n_node, n_edge, node_dim, edge_dim):
+    layer_dict = parse_GraphBlock(block_name, config, n_node, n_edge, node_dim, edge_dim)
+    layer_dict["class_name"] = "EdgeBlock"
+    layer_dict["inputs"] = [update_dict["last_node_update"], update_dict["last_edge_update"], "edge_index"]
+    layer_dict["outputs"] = [f"layer{index}_out"]
+    update_dict["last_edge_update"] = f"layer{index}_out"
+    return layer_dict, update_dict
+
+@pyg_handler('EdgeAggregate')
+def parse_EdgeAggregate(block_name, config, update_dict, index, n_node, n_edge, node_dim, edge_dim):
+    layer_dict = {"name": f"aggr{index}",
+                  "class_name": "EdgeAggregate",
+                  "n_node": n_node,
+                  "n_edge": n_edge,
+                  "node_dim": node_dim,
+                  "edge_dim": edge_dim,
+                  "out_dim": edge_dim,
+                  "inputs": [update_dict["last_edge_update"], "edge_index"],
+                  "outputs": [f"layer{index}_out"]}
+    update_dict["last_edge_aggr_update"] = f"layer{index}_out"
+    return layer_dict, update_dict

hls4ml/converters/pyg_to_hls.py

@@ -0,0 +1,149 @@
+from __future__ import print_function
+from collections import OrderedDict
+
+from hls4ml.converters.pytorch_to_hls import PyTorchModelReader
+from hls4ml.model.hls_model import HLSModel
+from hls4ml.templates import get_backend
+
+class PygModelReader(PyTorchModelReader):
+    def __init__(self, config):
+        super().__init__(config)
+        self.n_node = config['InputShape']['NodeAttr'][0]
+        self.n_edge = config['InputShape']['EdgeAttr'][0]
+        self.node_dim = config['InputShape']['NodeAttr'][1]
+        self.edge_dim = config['InputShape']['EdgeAttr'][1]
+
+    def get_weights_data(self, layer_name, var_name, module_name=None):
+        data = None
+
+        # Parameter mapping from pytorch to keras
+        torch_paramap = {
+            # Conv
+            'kernel': 'weight',
+            # Batchnorm
+            'gamma': 'weight',
+            'beta': 'bias',
+            'moving_mean': 'running_mean',
+            'moving_variance': 'running_var'}
+
+        if var_name not in list(torch_paramap.keys()) + ['weight', 'bias']:
+            raise Exception('Pytorch parameter not yet supported!')
+
+        if module_name is not None:
+            if var_name in list(torch_paramap.keys()):
+                var_name = torch_paramap[var_name]
+
+            try:
+                data = self.state_dict[module_name + '.' + layer_name + '.' + var_name].numpy().transpose()
+            except KeyError:
+                data = self.state_dict[module_name + '.layers.' + layer_name + '.' + var_name].numpy().transpose()
+
+        else:
+            if var_name in list(torch_paramap.keys()):
+                var_name = torch_paramap[var_name]
+
+            data = self.state_dict[layer_name + '.' + var_name].numpy().transpose()  # Look at transpose when systhesis produce lousy results. Might need to remove it.
+
+        return data
+
+# EdgeBlock/NodeBlock/Aggregate handlers
+block_handlers = {}
+
+def register_pyg_block_handler(block_name, handler_func):
+    if block_name in block_handlers:
+        raise Exception('Block {} already registered'.format(block_name))
+    else:
+        block_handlers[block_name] = handler_func
+
+def get_supported_pyg_blocks():
+    return list(block_handlers.keys())
+
+def pyg_handler(*args):
+    def decorator(function):
+        function.handles = [arg for arg in args]
+        return function
+    return decorator
+
+def pyg_to_hls(config):
+    forward_dict = config['ForwardDictionary']
+    activate_final = config['ActivateFinal']
+
+    # get precisions
+    backend = get_backend(config.get('Backend', 'Vivado'))
+    fp_type = backend.convert_precision_string(config['HLSConfig']['Model']['Precision'])
+    int_type = backend.convert_precision_string(config['HLSConfig']['Model']['IndexPrecision'])
+
+    # make reader
+    reader = PygModelReader(config)
+    n_node = reader.n_node
+    n_edge = reader.n_edge
+    node_dim = reader.node_dim
+    edge_dim = reader.edge_dim
+
+    # initiate layer list with inputs: node_attr, edge_attr, edge_index
+    layer_list = []
+    input_shapes = reader.input_shape
+    NodeAttr_layer = {
+        'name': 'node_attr',
+        'class_name': 'InputLayer',
+        'input_shape': input_shapes['NodeAttr'],
+        'inputs': 'input',
+        'dim_names': ['N_NODE', 'NODE_DIM'],
+        'precision': fp_type
+    }
+    layer_list.append(NodeAttr_layer)
+    EdgeAttr_layer = {
+        'name': 'edge_attr',
+        'class_name': 'InputLayer',
+        'input_shape': input_shapes['EdgeAttr'],
+        'inputs': 'input',
+        'dim_names': ['N_EDGE', 'EDGE_DIM'],
+        'precision': fp_type
+    }
+    layer_list.append(EdgeAttr_layer)
+    EdgeIndex_layer = {
+        'name': 'edge_index',
+        'class_name': 'InputLayer',
+        'input_shape': input_shapes['EdgeIndex'],
+        'inputs': 'input',
+        'dim_names': ['N_EDGE', 'TWO'],
+        'precision': int_type
+    }
+    layer_list.append(EdgeIndex_layer)
+    update_dict = {"last_node_update": "node_attr", "last_edge_update": "edge_attr", "last_edge_aggr_update": None}
+
+    # insert an aggregation step before each NodeBlock
+    aggr_count = 0
+    forward_dict_new = OrderedDict()
+    for key, val in forward_dict.items():
+        if val == "NodeBlock":
+            aggr_count += 1
+            aggr_key = f"aggr{aggr_count}"
+            aggr_val = "EdgeAggregate"
+            forward_dict_new[aggr_key] = aggr_val
+        forward_dict_new[key] = val
+
+    # complete the layer list
+    for i, (key, val) in enumerate(forward_dict_new.items()):
+        # get inputs, outputs
+        index = len(layer_list)+1
+        layer_dict, update_dict = block_handlers[val](key, config, update_dict, index, n_node, n_edge, node_dim, edge_dim)
+        layer_list.append(layer_dict)
+
+    if activate_final is not None:
+        act_dict = {
+            'name': 'final_act',
+            'class_name': 'Activation',
+            'inputs': [f"layer{len(layer_list)}_out"],
+            'activation': activate_final,
+            'precision': fp_type
+        }
+        layer_list.append(act_dict)
+        out = ["final_act"]
+    else:
+        out = [layer_list[-1]['outputs'][0]]
+
+    hls_model = HLSModel(config, reader, layer_list, inputs=['node_attr', 'edge_attr', 'edge_index'])
+    hls_model.outputs = out
+    return hls_model
+

hls4ml/converters/pytorch_to_hls.py

@@ -119,7 +119,7 @@ def pytorch_to_hls(config):
     -----
     Only sequential pytorch models are supported for now.
     """
-    
+
     #This is a list of dictionaries to hold all the layer info we need to generate HLS
     layer_list = []