def test_onnxruntime(model_name):
"""this function takes in a pytorch_model as input
"""
input_tensor = generate_test_input("pytorch", model_name)
torch_path, onnx_path = pytorch_onnx_paths(model_name)
torch_device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
torch_model = torch_load(torch_path, torch_device)
torch_model.eval()
torch_output = torch_model(input_tensor)
torch_onnx_export(torch_model,
input_tensor,
onnx_path,
export_params=True,
opset_version=9,
do_constant_folding=True,
input_names=['input'],
output_names=['output'],
dynamic_axes=None)
ort_outs = onnx_runtime(input_tensor, onnx_path, torch_output)
print(f"torch output: {torch_output}")
print(f"onnx output: {ort_outs}")
def main(model_name):
prebuilt_models = [
"resnet18", "alexnet", "vgg16", "squeezenet", "densenet", "inception",
"googlenet", "shufflenet", "mobilenet", "resnext50_32x4d",
"wide_resnet50_2", "mnasnet"
]
if model_name in prebuilt_models:
moodel = prebuilt_generator(model_name)
else:
model = TestNet()
model_path, _ = pytorch_onnx_paths(model_name)
save(model, model_path)
def torch_to_onnx(
model_name:str,
num_frames:int,
use_state_dict:bool,
return_models:bool=False)->None:
"""
Arg:
model_name (str): filename of the model
num_frames (int): number of feature frames that will fix the model's size
return_models (bool, False): if true, the function will return the torch and onnx model objects
"""
torch_path, config_path, onnx_path = pytorch_onnx_paths(model_name)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
config = load_config(config_path)
model_cfg = config['model']
freq_dim = 257 #freq dimension out of log_spectrogram
vocab_size = 39
time_dim = num_frames
model_cfg.update({'blank_idx': config['preproc']['blank_idx']})
torch_model = CTC_model(freq_dim, vocab_size, model_cfg)
state_dict = load_state_dict(torch_path, device=device)
torch_model.load_state_dict(state_dict)
torch_model.to(device)
print("model on cuda?: ", torch_model.is_cuda)
torch_model.eval()
# create the tracking inputs
hidden_size = config['model']['encoder']['rnn']['dim']
input_tensor = generate_test_input("pytorch", model_name, time_dim, hidden_size)
# export the models to onnx
torch_onnx_export(torch_model, input_tensor, onnx_path)
print(f"Torch model sucessfully converted to Onnx at {onnx_path}")
if return_models:
onnx_model = onnx.load(onnx_path)
return torch_model, onnx_model
def torch_to_coreml(
model_name:str,
num_frames:int,
use_state_dict:bool,
return_models:bool=False)->None:
"""
Arg:
model_name (str): filename of the model
num_frames (int): number of feature frames that will fix the model's size
return_models (bool, False): if true, the function will return the torch and onnx model objects
"""
torch_path, config_path, onnx_path = pytorch_onnx_paths(model_name)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
config = load_config(config_path)
model_cfg = config['model']
freq_dim = 257 #freq dimension out of log_spectrogram
vocab_size = 39
torch_model = CTC_model(freq_dim, vocab_size, model_cfg)
state_dict = load_state_dict(torch_path, device=device)
torch_model.load_state_dict(state_dict)
torch_model.to(device)
print("model on cuda?: ", torch_model.is_cuda)
torch_model.eval()
# create the tracking inputs
hidden_size = config['model']['encoder']['rnn']['dim']
x, (h_in, c_in) = generate_test_input("pytorch", model_name, 31, hidden_size)
traced_model = torch.jit.trace(torch_model, (x, (h_in, c_in)))
x_46, (h_46, c_46) = generate_test_input("pytorch", model_name, 46, hidden_size)
out_46, (h_out_46, c_out_46) = traced_model(x_46, (h_46, c_46))
if return_models:
pass
def test_onnxruntime(model_name):
"""this function takes in a pytorch_model as input
"""
#input_tensor = torch.randn(5, 3, 10, requires_grad=True).cuda()
#input_tensor = generate_test_input("pytorch", model_name)
input_tensor, h0, c0 = generate_test_input("pytorch", model_name)
torch_path, onnx_path = pytorch_onnx_paths(model_name)
torch_device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
torch_model = torch_load(torch_path, torch_device)
torch_model.eval()
#torch_output = torch_model(input_tensor)
torch_output, (hn, cn) = torch_model(input_tensor, (h0, c0))
print(f"torch_output: {torch_output}")
#torch_output = torch_export_inference(torch_path, onnx_path, input_tensor)
onnx_runtime(input_tensor, onnx_path, torch_output)