Unity-Technologies / barracuda-release

Other
564 stars 76 forks source link

Barracuda with savedmodel file format #326

Open romygt opened 1 year ago

romygt commented 1 year ago

Can I use barracuda with the savedModel format instead of Onnx

dilne commented 1 year ago

ONNX is required. You can convert a PyTorch saved model to ONNX like shown below. I include the ONNX checker, it should confirm everything worked:

import torch
import torch.onnx
import onnx
import numpy as np

# load path
load_path = r"your/model/directory/name_of_model.pt"

# output path
output_path = r"your/output/directory/name_of_model.onnx"

model = torch.load(load_path)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
model.eval()

# Input to the model (you need to define the size for the model inputs)
# This example is for a single batch, single channel tensor of size 28 x 28)
x = torch.randn(1, 1, 28, 28, requires_grad=True).to(device)
torch_out = model(x)

# Export the model
torch.onnx.export(model,               # model being run
                  x,                         # model input (or a tuple for multiple inputs)
                  output_path,               # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=9,           # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names = ['input'],   # the model's input names
                  output_names = ['output'], # the model's output names
                  )

onnx_model = onnx.load(output_path)
onnx.checker.check_model(onnx_model)
hayden-donnelly commented 1 year ago

Not sure if Tensorflow savedmodel is any different, but I use this code to convert my TF savedmodels to ONNX:

import tensorflow as tf 
import tf2onnx

input_path = "../data/models/"
model_name = "simple_upsampler_bilinear"
file_type = ''
output_path = "../data/models/onnx_models/"

#load the model.
pre_model = tf.keras.models.load_model(input_path + model_name + file_type)

# Convert h5 to onnx.
tf2onnx.convert.from_keras(pre_model, output_path = output_path + model_name + ".onnx", opset = 9)