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)