3. Importing graphs

The PopART Session class creates the runtime environment for executing graphs on IPU hardware. It can read an ONNX graph from a serialised ONNX model protobuf (ModelProto), either directly from a file or from memory. A Session object can be constructed either as InferenceSession (Python, C++) for inference or TrainingSession (Python, C++) for training.

Some metadata must be supplied to construct the Session class. These are described in Section 3.1, Creating a session.

In the following example of importing a graph for inference, the torchvision package is used to create a pre-trained AlexNet graph , with a 4 x 3 x 244 x 244 input. The graph has an ONNX output called output, and the DataFlow object (Python, C++) contains an entry to fetch that anchor tensor.

# Copyright (c) 2020 Graphcore Ltd. All rights reserved.
import popart

import torch.onnx
import torchvision

input_ = torch.FloatTensor(torch.randn(4, 3, 224, 224))
model = torchvision.models.alexnet(pretrained=False)

output_name = "output"


def onnx_opset_version():
    popart_max_onnx_opset_version = 11
    default_onnx_opset_version = 0

    try:
        from torch.onnx.symbolic_helper import _export_onnx_opset_version
    except ImportError:
        try:
            from torch.onnx._globals import GLOBALS
        except ImportError:
            default_onnx_opset_version = popart_max_onnx_opset_version
        else:
            default_onnx_opset_version = GLOBALS.export_onnx_opset_version
    else:
        default_onnx_opset_version = _export_onnx_opset_version

    return min(popart_max_onnx_opset_version, default_onnx_opset_version)


torch.onnx.export(
    model,
    input_,
    "alexnet.onnx",
    output_names=[output_name],
    opset_version=onnx_opset_version(),
)

# Create a runtime environment
anchors = {output_name: popart.AnchorReturnType("All")}
dataFlow = popart.DataFlow(100, anchors)
device = popart.DeviceManager().createCpuDevice()

session = popart.InferenceSession("alexnet.onnx", dataFlow, device)

files/importing_graphs.py

The DataFlow object is described in more detail in Section 5, Executing graphs.

3.1. Creating a session

The following parameters are required to create a Session object:

  • model: The name of a protobuf file, or the protobuf itself in order to get the ONNX graph.

  • dataFlow: A DataFlow object which contains the following needed to execute the graph:

    • Batches per step: The number of batches to run in a single call to Session::run():

      • For InferenceSession, this is the number of executions of the model.

      • For TrainingSession, this is the number of weight updates.

    • The names of the tensors in the graph used to return the results to the host.

  • inputShapeInfo: In some ONNX graphs, the sizes of input tensors might not be specified. In this case, the optional inputShapeInfo parameter can be used to specify the input shapes. The Poplar framework uses statically allocated memory buffers and so it needs to know the size of tensors before the compilation.

  • patterns: The optional patterns parameter allows the user to select a set of graph transformation patterns which will be applied to the graph. Without this parameter, a default set of optimisation transformations will be applied.

  • userOptions: The options to be applied to the session. This is described in more detail in Section 3.2, Session control options.

  • For TrainingSession only:

    • loss: The types of loss to apply to the network.

    • optimizer: The optimiser to use.

An example of creating a Session object from an ONNX model is shown below.

# Copyright (c) 2020 Graphcore Ltd. All rights reserved.
import popart

import torch.onnx
import torchvision

input_ = torch.FloatTensor(torch.randn(4, 3, 224, 224))
model = torchvision.models.alexnet(pretrained=False)

output_name = "output"


def onnx_opset_version():
    popart_max_onnx_opset_version = 11
    default_onnx_opset_version = 0

    try:
        from torch.onnx.symbolic_helper import _export_onnx_opset_version
    except ImportError:
        try:
            from torch.onnx._globals import GLOBALS
        except ImportError:
            default_onnx_opset_version = popart_max_onnx_opset_version
        else:
            default_onnx_opset_version = GLOBALS.export_onnx_opset_version
    else:
        default_onnx_opset_version = _export_onnx_opset_version

    return min(popart_max_onnx_opset_version, default_onnx_opset_version)


torch.onnx.export(
    model,
    input_,
    "alexnet.onnx",
    output_names=[output_name],
    opset_version=onnx_opset_version(),
)

# Create a runtime environment
anchors = {output_name: popart.AnchorReturnType("All")}
dataFlow = popart.DataFlow(100, anchors)

# Append an Nll loss operation to the model
builder = popart.Builder("alexnet.onnx")
labels = builder.addInputTensor("INT32", [4], "label_input")
nlll = builder.aiGraphcore.nllloss([output_name, labels])

optimizer = popart.ConstSGD(0.001)

# Run session on CPU
device = popart.DeviceManager().createCpuDevice()
session = popart.TrainingSession(
    builder.getModelProto(),
    deviceInfo=device,
    dataFlow=dataFlow,
    loss=nlll,
    optimizer=optimizer,
)

files/importing_session.py

In this example, when the TrainingSession object is created, a negative log likelihood (NLL) loss node (Python, C++) will be added to the end of the graph, and a ConstSGD optimiser will be used to optimise the parameters in the network.

3.2. Session control options

The optional userOptions parameter passes options to the session that control specific features of the PopART session. The available PopART options are listed in Section 5.2, Session options in the PopART C++ API reference.

The userOptions parameter also controls the underlying Poplar functions:

  • engineOptions passes options to the Poplar Engine object created to run the graph.

  • convolutionOptions passes options to the PopLibs convolution functions.

  • reportOptions controls the instrumentation and generation of profiling information.

Full details of the Poplar options can be found in the Poplar and PopLibs API Reference.

Section 5.6, Retrieving profiling reports contains examples of how to use some of these options.

3.3. Executing an imported graph

Now that the device has been selected, the graph can be compiled for it and loaded onto the hardware. The prepareDevice() method (Python, C++) allows you to do this.

In order to execute the graph, input data has to be provided. The input created here to feed data to PopART must be a NumPy array, rather than an initialised torch tensor.

Finally, the PyStepIO class (Python, IStepIO in C++) provides a session with input and output buffers. For both input and output, this class takes a dictionary with tensor names as keys and Python (or NumPy) arrays as values. Note that, for the imported graph, the key should match the tensor names in the graph. In this case, input.1 is the input tensor name in the imported graph alexnet.onnx, and input_1 is the value fed into it.

In order to find the input names, you can import the onnx package and use onnx.load to load the model. loaded_model.graph.input and loaded_model.graph.output give you all the node information for inputs and outputs. You can use the following command to extract the names of inputs and outputs:

inputs_name = [node.name for node in loaded_model.graph.input]
outputs_name = [node.name for node in loaded_model.graph.output]

An example of executing an imported graph is shown below.

# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
import popart

import torch.onnx
import torchvision

import numpy as np
import onnx

input_ = torch.FloatTensor(torch.randn(4, 3, 224, 224))
model = torchvision.models.alexnet(pretrained=False)

output_name = "output"


def onnx_opset_version():
    popart_max_onnx_opset_version = 11
    default_onnx_opset_version = 0

    try:
        from torch.onnx.symbolic_helper import _export_onnx_opset_version
    except ImportError:
        try:
            from torch.onnx._globals import GLOBALS
        except ImportError:
            default_onnx_opset_version = popart_max_onnx_opset_version
        else:
            default_onnx_opset_version = GLOBALS.export_onnx_opset_version
    else:
        default_onnx_opset_version = _export_onnx_opset_version

    return min(popart_max_onnx_opset_version, default_onnx_opset_version)


torch.onnx.export(
    model,
    input_,
    "alexnet.onnx",
    output_names=[output_name],
    opset_version=onnx_opset_version(),
)

# Obtain inputs/outputs name of loaded model
loaded_model = onnx.load("alexnet.onnx")
inputs_name = [node.name for node in loaded_model.graph.input]
outputs_name = [node.name for node in loaded_model.graph.output]

print("Iputs name:", inputs_name)
print("Outputs name:", outputs_name)

# Create a runtime environment
anchors = {output_name: popart.AnchorReturnType("All")}
dataFlow = popart.DataFlow(100, anchors)
device = popart.DeviceManager().createCpuDevice()

session = popart.InferenceSession("alexnet.onnx", dataFlow, device)

session.prepareDevice()

input_1 = np.random.randn(4, 3, 224, 224).astype(np.float32)

stepio = popart.PyStepIO({"input.1": input_1}, session.initAnchorArrays())