def test_customized_model_multiple_inputs_and_outputs():
# prepare model
model_path = os.path.join(TEST_INPUT_DIR,
'customized_model_multiple_inputs_outputs',
f"_{str(uuid.uuid4())}.pth")
inputs_schema = [{
IOSchemaLoader.SHAPE_KEY: [10, 3, 244, 244],
IOSchemaLoader.DTYPE_KEY: "float",
IOSchemaLoader.NAME_KEY: "multiple_in_0"
}, {
IOSchemaLoader.SHAPE_KEY: [10, 3, 244, 244],
IOSchemaLoader.DTYPE_KEY: "float",
IOSchemaLoader.NAME_KEY: "multiple_in_1"
}]
outputs_schema = [{
IOSchemaLoader.NAME_KEY: "multiple_out_0"
}, {
IOSchemaLoader.NAME_KEY: "multiple_out_1"
}]
onnx_model_path = f"_{str(uuid.uuid4())}.onnx"
prepare_model(model_path, "customized_model_multiple_inputs_outputs",
len(inputs_schema))
cvt_config = ConversionConfig(
model_path=model_path,
inputs_schema=inputs_schema,
outputs_schema=outputs_schema,
sample_input_data_path=CUSTOMIZED_MODEL_MULTIPLE_INPUTS_OUTPUTS_DATA,
model_framework="pytorch",
onnx_model_path=onnx_model_path)
convert(cvt_config)
os.remove(onnx_model_path)
def test_pretrained_model_classification(model_name):
# prepare model
model_path = os.path.join(TEST_INPUT_DIR, model_name,
f"_{str(uuid.uuid4())}.pth")
prepare_model(model_path, model_name)
# test
inputs_schema = [{
IOSchemaLoader.SHAPE_KEY: [1, 3, 244, 244],
IOSchemaLoader.DTYPE_KEY: "float",
IOSchemaLoader.NAME_KEY: "input_0"
}]
outputs_schema = [{IOSchemaLoader.NAME_KEY: "output_0"}]
onnx_model_path = f"_{str(uuid.uuid4())}.onnx"
cvt_config = ConversionConfig(model_path=model_path,
inputs_schema=inputs_schema,
outputs_schema=outputs_schema,
model_framework="pytorch",
onnx_model_path=onnx_model_path,
onnx_opset=11)
convert(cvt_config)
os.remove(onnx_model_path)
def test_pretrained_model_video_r3d_18():
# prepare model
model_path = os.path.join(TEST_INPUT_DIR, 'video_r3d_18',
f"_{str(uuid.uuid4())}.pth")
prepare_model(model_path, "video_r3d_18", 1)
# test
inputs_schema = [{
IOSchemaLoader.SHAPE_KEY: [2, 3, 4, 112, 112],
IOSchemaLoader.DTYPE_KEY: "float",
IOSchemaLoader.NAME_KEY: "input_0"
}]
outputs_schema = [{IOSchemaLoader.NAME_KEY: "output_0"}]
onnx_model_path = f"_{str(uuid.uuid4())}.onnx"
cvt_config = ConversionConfig(
model_path=model_path,
inputs_schema=inputs_schema,
outputs_schema=outputs_schema,
sample_input_data_path=PRETRAINED_MODEL_VIDEO_DATA,
onnx_opset=11,
model_framework="pytorch",
onnx_model_path=onnx_model_path)
convert(cvt_config)
os.remove(onnx_model_path)