Model API Outline

To send a model to be displayed in the Zetane Engine, you need to add the following few lines to your ML script. Depending on your framework, you can take the appropriate the ONNX, Keras or Pytorch line below.

import zetane

# Launch the Zetane Engine
engine = ztn.Context()

# Create the model object to send to the engine
zmodel = engine.model()

# ONNX: pass the ONNX file path and inputs to the zmodel
zmodel.onnx(onnx_path).inputs(input_path)
# Or
# Keras: pass the Keras model and inputs to the zmodel
zmodel.keras(model).inputs(inputs)
# Or
# Pytorch: pass the Pytorch model and inputs to the zmodel
zmodel.torch(model).inputs(inputs)

# Update the model and send to the engine
zmodel.update()

# Disconnect from the engine
engine.disconnect()

If wanted, it is possible to do function chaining and reduce the number of lines needed. For example, we can add the following lines to Pytorch script.

import zetane
engine = ztn.Context()

zmodel = engine.model().torch(model).inputs(inputs).update()

engine.disconnect()

Hello Keras MNIST

Sending your model from a script and display it in the engine is very similar to sending dynamic text as seen in the previous section. We started with a basic MNIST example script which can be found the on the Keras website and inserted code display in the engine the model, input and output to the engine.

There are two main parts to display your model in the Zetane engine.

Create a model object:

zmodel = engine.model().keras(model)

Update the model object with an input of your choice:

zmodel.update(inputs=x_exp)

Overview of the code block insertions:

  • Import zetane

  • Created Zetane objects

  • Build a loop taking some test data to make predictions

  • Send or sync the input, model and output so that they can be displayed

The blocks of code that were inserted to the Keras MNIST sample script are found under the # ZETANE API and # PREDICT comments.

# Adapted from: https://keras.io/examples/vision/mnist_convnet
# IMPORT MODULES
import os
import sys
import time
import numpy as np
from tensorflow import keras
from tensorflow.keras import layers
import tensorflow as tf

# ZETANE API: import module
import zetane

# DATA PREPARATION
# Model / data parameters
num_classes = 10
input_shape = (28, 28, 1)

# the data, split between train and test sets
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

# Scale images to the [0, 1] range
x_train = x_train.astype("float32") / 255
x_test = x_test.astype("float32") / 255
# Make sure images have shape (28, 28, 1)
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)
print("x_train shape:", x_train.shape)
print(x_train.shape[0], "train samples")
print(x_test.shape[0], "test samples")

# convert class vectors to binary class matrices
y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)

# BUILD THE MODEL
model = keras.Sequential(
        [
                keras.Input(shape=input_shape),
                layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
                layers.MaxPooling2D(pool_size=(2, 2)),
                layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
                layers.MaxPooling2D(pool_size=(2, 2)),
                layers.Flatten(),
                layers.Dropout(0.5),
                layers.Dense(num_classes, activation="softmax"),
        ]
)

# TRAIN THE MODEL
model.summary()
batch_size = 128
epochs = 1
model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])
model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_split=0.1)

# ZETANE API: create the objects
engine = zetane.Context()
zmodel = engine.model().keras(model)
zimg = engine.image().position(-1.0, 1.5, 0.0)
ztxt_target = engine.text().position(-1.5, -1, 1).font_size(0.16)
ztxt_pred = engine.text().position(-1.5, -1.5, 1).font_size(0.16)

# PREDICT WITH THE TRAIN MODEL AND SEND TO THE ENGINE
test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))
# perform inference for elements of test dataset
for x, y in test_dataset.take(1000):
        target = tf.argmax(y).numpy()

        x_exp = tf.expand_dims(x, 0).numpy()
        y_out = model.predict(x_exp, steps=1)
        pred = tf.argmax(y_out[0]).numpy()

        # ZETANE API: update the objects to inspect them in the engine
        # update input image
        zimg.update(data=x.numpy())
        # update model and internal tensors
        zmodel.update(inputs=x_exp)
        # update output
        ztxt_target.text("Target class: " + str(target)).update()
        ztxt_pred.text("Predicted class: " + str(pred)).update()

        time.sleep(0.1)

Hello PyTorch MNIST

Sending a PyTorch model from a script and displaying it in the engine is very similar to sending a Keras model as seen in the previous section. We presented a basic MNIST example script to display the model, input and output in the engine.

There are two main parts to display your PyTorch model in the Zetane engine.

Create a model object:

zmodel = engine.model().torch(model, inputs)

Note that in the case of PyTorch, the model input (real or fake, but matching the dimensions) should be passed when initializing the model.

Update the model object with an input of your choice:

zmodel.update(inputs=x_exp)

Overview of the code block insertions:

  • Import zetane

  • Created Zetane objects

  • Build a loop taking some test data to make predictions

  • Send or sync the input, model and output so that they can be displayed

# Adapted from: https://pytorch.org/tutorials/beginner
# IMPORT MODULES
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import zetane
import time
import matplotlib.pyplot as plt
import numpy as np
import torchvision

# DEFINE THE MODEL
class Net(nn.Module):
        def __init__(self):
                super(Net, self).__init__()
                self.conv1 = nn.Conv2d(1, 32, 3, 1)
                self.conv2 = nn.Conv2d(32, 64, 3, 1)
                self.conv3 = nn.Conv2d(64, 64, 3, 1)
                self.conv4 = nn.Conv2d(64, 16, 3, 1)
                self.dropout = nn.Dropout(0.5)
                self.fc = nn.Linear(1024, 10)
                self.pool = nn.MaxPool2d(2, 2)

        def forward(self, x):
                x = self.conv1(x)
                x = F.relu(x)
                x = self.conv2(x)
                x = F.relu(x)
                x = self.pool(x)
                x = self.conv3(x)
                x = F.relu(x)
                x = self.conv4(x)
                x = F.relu(x)
                x = torch.flatten(x, 1)
                x = self.dropout(x)
                x = self.fc(x)
                output = F.log_softmax(x, dim=1)
                return output

net = Net()
#print(net)

def imshow(img):
        img = img * 0.3081 + 0.1307
        npimg = img.numpy()
        print(img.shape)
        plt.imshow(np.transpose(npimg, (1, 2, 0)))
        plt.show()


classes=range(10)

# LOSS FUNCTION AND OPTIMIZER
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

# TRAINING AND TEST DATA
transform=transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize((0.1307,), (0.3081,))
                ])
trainset = datasets.MNIST(root='./data', train=True,
                                                                                download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                                                                  shuffle=True, num_workers=0)

testset = datasets.MNIST(root='./data', train=False,
                                                                           download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=1,
                                                                                 shuffle=False, num_workers=0)

# ZETANE API
engine = zetane.Context()
engine.clear_universe()
zmodel = engine.model()
zimg = engine.image().position(-1.0, 1.5, 0.0)
ztxt_target = engine.text().position(-1.5, -1, 1).font_size(0.16)
ztxt_pred = engine.text().position(-1.5, -3, 1).font_size(0.16)

#TRAIN THE MODEL
num_epoch = 1
print('Started Training for {} epoch(s)'.format(num_epoch))
for epoch in range(num_epoch):  # loop over the dataset multiple times
        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
                # get the inputs; data is a list of [inputs, labels]
                inputs, labels = data
                # zero the parameter gradients
                optimizer.zero_grad()

                # forward + backward + optimize
                outputs = net(inputs)
                loss = criterion(outputs, labels)
                loss.backward()
                optimizer.step()

                # print statistics
                running_loss += loss.item()
                if i % 2000 == 1999:    # print every 2000 mini-batches
                        print('[%d, %5d] loss: %.3f' %
                                  (epoch + 1, i + 1, running_loss / 2000))
                        running_loss = 0.0

print('Finished Training')


# PREDICT WITH THE TRAIN MODEL AND SEND TO THE ENGINE
dataiter = iter(testloader)
images, labels = next(dataiter)
zmodel.torch(net, images).update()

for itr in range(100):
        images, labels = next(dataiter)
        print('GroundTruth: {}'.format(classes[labels[0]]))
        outputs = net(images)
        _, predicted = torch.max(outputs, 1)

        print('Predicted: {}'.format(classes[predicted[0]]))


        # ZETANE API: update the objects to inspect them in the engine
        # update input image
        zimg.update(data=images[0,0,:,:].numpy())
        # update model and internal tensors
        zmodel.update(inputs=np.array(images))
        # update output
        ztxt_target.text("Target class: " + str(classes[labels[0]])).update()
        ztxt_pred.text("Predicted class: " + str(classes[predicted[0]])).update()

        time.sleep(0.1)

Hello ONNX Runtime

This is a basic example to send and visualize ONNX models in the engine. For most models of the ONNX Model Zoo, a similar script would work.

Before running the script, make sure you have downloaded the pre-trained .ONNX model and the test image test.jpg in your folder.

Download Pre-trained ONNX model

Download the test image (test.jpg)

# Original script: ONNX Model Zoo, Super Resolution
# https://github.com/onnx/models/tree/master/vision/super_resolution/sub_pixel_cnn_2016

from PIL import Image
import numpy as np
import onnxruntime as rt
import os
import sys
import zetane as ztn

# Launch Zetane
zcontext = ztn.Context()
zcontext.clear_universe()
zonnx = zcontext.model()

zimg_input = zcontext.image()
zimg_output = zcontext.image()

# provide the name of the model and width and height of the images for the variables
model = 'super_resolution.onnx'
width = 224
height = 224

# Load and downscale image
input_img_path = 'test.jpg'
orig_img = Image.open(input_img_path)
downscaled_img = orig_img.resize((width, height))
downscaled_img.save("test_downscaled.jpg")

# Preprocess image for the model
img_ycbcr = downscaled_img.convert('YCbCr')
img_y_0, img_cb, img_cr = img_ycbcr.split()
img_ndarray = np.asarray(img_y_0)
img_4 = np.expand_dims(np.expand_dims(img_ndarray, axis=0), axis=0)
img_5 = img_4.astype(np.float32) / 255.0
img_5

# Load onnx file and run inference
session = rt.InferenceSession('super_resolution.onnx')
output_name = session.get_outputs()[0].name
input_name = session.get_inputs()[0].name
result = session.run([output_name], {input_name: img_5})
img_out_y = result[0]
print(img_out_y.shape)

# Postprocess
img_out_y = Image.fromarray(np.uint8((img_out_y[0] * 255.0).clip(0, 255)[0]), mode='L')
# get the output image following the post-processing step from the PyTorch implementation
final_img = Image.merge(
                "YCbCr", [
                                img_out_y,
                                img_cb.resize(img_out_y.size, Image.BICUBIC),
                                img_cr.resize(img_out_y.size, Image.BICUBIC),
                ]).convert("RGB")

# Display in the Zetane engine
# Downscaled image
image_for_zetane = np.asarray(downscaled_img)
image_for_zetane = image_for_zetane.astype(np.float32) / 255.0
zimg_input.position(-3.5, 1, 0).scale(0.1, 0.1, 0.1).update(data=image_for_zetane)

# Model architecture and tensors
zonnx.model(model).update(img_5)

# Model output super resolution image
final_img = np.asarray(final_img)
final_img = final_img.astype(np.float32) / 255.0
zimg_output.position(-0.5, 1, 0).scale(0.0335,0.0335,0.0335).update(data=final_img)