Aidge tiling demonstration

This tutorial aims at demonstrating how tiling can splite computation on several devices.

import aidge_core
import aidge_backend_cpu
import aidge_onnx
import numpy as np

Define mermaid visualizer function

Aidge save graph using the mermaid format, in order to visualize the graph live in the notebook, we will setup the following function:

import base64
from IPython.display import Image, display
import matplotlib.pyplot as plt

def visualize_mmd(path_to_mmd):
    with open(path_to_mmd, "r") as file_mmd:
       graph_mmd = file_mmd.read()

    graphbytes = graph_mmd.encode("utf-8")
    base64_bytes = base64.b64encode(graphbytes)
    base64_string = base64_bytes.decode("utf-8")
    display(Image(url=f"https://mermaid.ink/img/{base64_string}"))

Let’s create a small neural network with four layers.

The GraphView generating function sequential is used. You should at least name the layers of most interest to ease access to them if required.

model = aidge_core.sequential([
                    aidge_core.LeakyReLU(1, name="leakyrelu0"),
                    aidge_core.Conv2D(3, 32, [3, 3], name="conv0"),
                    aidge_core.BatchNorm2D(32, name="bn0"),
                    aidge_core.ReLU(name="relu0")
                ])
model.save("initial_graph")

visualize_mmd("initial_graph.mmd")

Let’s create an input to link to the model.

# Create an input
input_tensor = aidge_core.Tensor(np.random.rand(4, 3, 66, 66).astype(np.float32))

Let’s generate random values for each parameter

convW = aidge_core.Tensor(np.random.rand(32, 3, 3, 3).astype(np.float32))
convB = aidge_core.Tensor(np.random.rand(32).astype(np.float32))
BNscale = aidge_core.Tensor(np.random.rand(32).astype(np.float32))
BNshift = aidge_core.Tensor(np.random.rand(32).astype(np.float32))
BNmean = aidge_core.Tensor(np.random.rand(32).astype(np.float32))
BNvar = aidge_core.Tensor(np.random.rand(32).astype(np.float32))

model.get_node("conv0").get_operator().set_input(1, convW)
model.get_node("conv0").get_operator().set_input(2, convB)

model.get_node("bn0").get_operator().set_input(1, BNscale)
model.get_node("bn0").get_operator().set_input(2, BNshift)
model.get_node("bn0").get_operator().set_input(3, BNmean)
model.get_node("bn0").get_operator().set_input(4, BNvar)

Select an implementation and compute input/output dimensions.

model.compile("cpu", aidge_core.dtype.float32, dims=[[4,3,66,66]])

Run the model

# Create SCHEDULER
scheduler = aidge_core.SequentialScheduler(model)

# Run inference !
scheduler.forward(data=[input_tensor])

# keep result in memory
res1 = np.array(model.get_node("relu0").get_operator().get_output(0))

Thanks to tiling, the convolution computation can be divided in the desired number of stripes.

Here, we choose 4 stripes on the second axis (the horizontl axis).

tiled_conv = aidge_core.get_conv_horizontal_tiling(model.get_node("conv0"), 2, 4)
node_to_replace = {model.get_node("conv0"),
                   model.get_node("conv0").get_parent(1),
                   model.get_node("conv0").get_parent(2)}

aidge_core.GraphView.replace(node_to_replace, tiled_conv)

True

The replace function returned True: the replacement was successful. We can visualize the new model.

The convolution has been divided in 4 smaller convolutions preceeded by a Slice operator to extract the right sub-tensor. All 4 results are concatenated back to a single Tensor that serves as an input for the following layer.

model.save("model_after_tiling")
visualize_mmd("model_after_tiling.mmd")

Now we run the transformed model and compare it’s ouput value to the previous one.

model.compile("cpu", aidge_core.dtype.float32)
scheduler.resetScheduling()
scheduler.forward(data=[input_tensor])
res2 = np.array(model.get_node("relu0").get_operator().get_output(0))

(res1 == res2).all()

True

Both outputs are the same !