How to load custom CNN model into Jetson Nano

Hello there,

So my objective is to load my custom CNN model into Jetson Nano and test it in real life.

What I’ve done:

  1. Successfully trained my CNN model on Google Colab Pro.
  2. Saved the hdf5 file from Google Colab with a size of 13MB~.
  3. Connected the Raspberry Pi V2 camera with Jetson Nano.

So here I am stuck in this process. I have gone through several tutorials on how to implement ImageNet using Jetson Nano, however, most of them are using the pre-trained models offered by NVIDIA such as (GoogleNet, AlexNet, ResNet, etc).

Is there any simple way/tutorial for me to start on the journey to load my custom CNN model file into Jetson Nano since this is the first time I’ve experienced it?

Thank you!

Try to use onnx format as Jeston Nano cannot withstand 16 MB model:

#Import libraries
import tensorflow as tf
from tensorflow.keras.models import load_model
import tf2onnx
import numpy as np
import onnxruntime as rt
import cv2

# Load model
model = load_model(“MODEL+PATH”)

#Convert model to onnx
spec = (tf.TensorSpec((1, 300, 300, 1), tf.float32, name=“input”),) #change image size as per your trained dataset
output_path = model.name + “.onnx”
model_proto, _ = tf2onnx.convert.from_keras(model, input_signature=spec, opset=13, output_path=output_path)
output_names = [n.name for n in model_proto.graph.output]
print(output_names)

#Load onnx model
image_shape = (300,300,1)
output_path = “ONNX_MODEL_PATH” #specify your ONNX converted model path
providers = [‘CPUExecutionProvider’]
m = rt.InferenceSession(output_path, providers=providers)

#Run prediction
img_path = “PATH_TO_IMAGE”
img_pred = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
print(img_pred.shape)
img_pred = cv2.resize(img_pred,(300,300))
img_pred = img_pred / 255 # rescale
onnx_pred = m.run(output_names, {“input”: img_pred.reshape(1,*image_shape).astype(np.float32)})

check the website of onnx for more details

1 Like

This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.