pythontensorflowpytorchtensorflow-xla

HLO protobuf to pytorch / tensorflow graph


Assume we have HLO protobuf from a model through Pytorch-XLA or Tensorflow.

  1. Is there a way to create computational graph from it?
  2. Is it possible to create Pytorch-XLA and Tensorflow model from it?

In python, input has following type [link].

hlo_pb2.HloModuleProto()

Solution

  • I think one approach would be to follow this two0step process.

    1. Extract relevant information from this protobuf object, such as the nodes, their connections, and attributes, and proceed with the conversion process into ONNX
    2. TensorFlow provides built-in utilities for converting ONNX models into TensorFlow models. After converting the HLO protobuf to ONNX format using the appropriate tools or libraries. Then, use TensorFlow's tf.compat.v1.graph_util.import_graph_def() or tf.saved_model.loader.load() to load the ONNX file and obtain a TensorFlow computational graph.

    Here is a sample snippet that you would need to adjust to achieve step 1

    Here's an example code snippet that demonstrates how to extract information from an HLO protobuf object and convert it into an ONNX file using PyTorch:

    import torch
    import onnx
    from onnx import helper
    from onnx import AttributeProto, TensorProto, GraphProto
    
    def convert_hlo_to_onnx(hlo_module):
        # Create an empty ONNX graph
        graph = helper.make_graph([], "hlo_model", [], [])
    
        # Keep track of the ONNX node names and their corresponding outputs
        node_outputs = {}
    
        # Iterate over the HLO computations and instructions
        for computation in hlo_module.computations:
            # Iterate over the HLO instructions in the computation
            for instruction in computation.instructions:
                # Extract instruction attributes
                instruction_name = instruction.name
                instruction_opcode = instruction.opcode
                instruction_outputs = instruction.operand
    
                # Create ONNX node with corresponding inputs and outputs
                onnx_node = helper.make_node(instruction_opcode, instruction.operand, instruction_outputs, name=instruction_name)
    
                # Set any additional attributes for the ONNX node if needed
                # For example, if the HLO instruction has attributes 'attr1' and 'attr2',
                # you can add them to the ONNX node as follows:
                # onnx_node.attribute.extend([
                #     helper.make_attribute('attr1', instruction.attr1),
                #     helper.make_attribute('attr2', instruction.attr2)
                # ])
    
                # Add the ONNX node to the graph
                graph.node.extend([onnx_node])
    
                # Update the node_outputs dictionary with the current ONNX node outputs
                node_outputs[instruction_name] = instruction_outputs
    
        # Iterate over the HLO computations again to establish connections between nodes
        for computation in hlo_module.computations:
            # Iterate over the HLO instructions in the computation
            for instruction in computation.instructions:
                # Get the current instruction name and its outputs
                instruction_name = instruction.name
                instruction_outputs = instruction.operand
    
                # Iterate over the outputs and connect them to the corresponding inputs
                for output in instruction_outputs:
                    # Check if the output is used as an input in any subsequent instruction
                    if output in node_outputs:
                        # Get the corresponding ONNX node output name
                        onnx_node_output = node_outputs[output]
    
                        # Find the ONNX node with the same name as the current instruction
                        onnx_node = next(node for node in graph.node if node.name == instruction_name)
    
                        # Update the inputs of the ONNX node to connect with the output of the previous node
                        onnx_node.input.remove(output)
                        onnx_node.input.extend(onnx_node_output)
    
        # Create the ONNX model with the graph
        model = helper.make_model(graph)
    
        # Save the ONNX model to a file
        onnx.save_model(model, "converted_model.onnx")
    
    # Assuming you have an HLO module protobuf object named 'hlo_module'
    hlo_module = hlo_pb2.HloModuleProto()
    
    # Call the conversion function
    convert_hlo_to_onnx(hlo_module)
    

    In this example, we ie iterate over the computations and instructions fields of the HloModuleProto. Each instruction represents an HLO instruction within a computation. We extract the instruction name, opcode, and outputs. Then, we create an ONNX node using helper.make_node() and set the inputs, outputs, and attributes accordingly. We add the ONNX node to the graph and keep track of the node outputs in the node_outputs dictionary.

    After creating the ONNX graph, we use helper.make_model() to create an ONNX model with the graph, and finally, onnx.save_model() is used to save the ONNX model to a file ("converted_model.onnx" in this example).

    Please note that this code assumes you have the necessary dependencies installed, including the onnx and torch packages. Also, make sure to import the relevant modules (onnx, helper, AttributeProto, TensorProto, GraphProto) from the onnx package.

    Remember to adapt this code to fit your specific HLO protobuf structure and attributes.