SO I have a model with input_name = [‘input_ids’, ‘style’,‘speed’] .
For clearly here is the ONNX creation Function
input_ids = torch.randint(1, 100, (48,)).numpy()
input_ids = torch.LongTensor([[0, *input_ids, 0]])
style = torch.randn(1, 256)
speed = torch.randint(1, 10, (1,)).int()
print("Input Ids = ",input_ids.shape)
print("Style = ",style.shape)
print("Speed = ",speed.shape)
torch.onnx.export(
model,
args = (input_ids, style, speed),
f = onnx_file,
export_params = True,
verbose = True,
input_names = [ 'input_ids', 'style', 'speed' ],
output_names = [ 'waveform', 'duration' ],
opset_version = 19,
dynamic_axes = {
'input_ids': { 1: 'input_ids_len' },
'waveform': { 0: 'num_samples' },
},
do_constant_folding = True,
)
print('export kokoro.onnx ok!')
Now here is My RKNN build function
if __name__ == '__main__':
model_path, platform, do_quant, output_path = "onnx/kokoro.onnx" , "rk3588" , False , "onnx/kokoro.rknn" #parse_arg()
input_shapes = {
'input_ids':[1,50],
'style':[1,256],
'speed':[1]
}
input_name = ['input_ids', 'style','speed']
input_size_list = [[1,50],[1,256],[1]]
dynamic_shapes = [
[[1, 50], [1, 256], [1]], # First possible combination
[[1, 100], [1, 256], [1]], # Second possible combination
[[1, 512], [1, 256], [1]] # Third possible combination
]
# Create RKNN object
rknn = RKNN(verbose=True)
# Pre-process config
print('--> Config model')
rknn.config(target_platform=platform,
quantized_dtype="w8a8"
,dynamic_input = dynamic_shapes
)
print('done')
# Load model
print('--> Loading model')
ret = rknn.load_onnx(model=model_path, input_size_list=input_size_list,
inputs=input_name,
)
if ret != 0:
print('Load model failed!')
exit(ret)
print('done')
# Build model
print('--> Building model')
ret = rknn.build(do_quantization=do_quant)
if ret != 0:
print('Build model failed!')
exit(ret)
print('done')
# Export rknn model
print('--> Export rknn model')
ret = rknn.export_rknn(output_path)
if ret != 0:
print('Export rknn model failed!')
exit(ret)
print('done')
# Release
rknn.release()
Still Getting error IN build stage