def test_multiple_inputs():
# prepare convert parameters
inputs_schema = []
input_names = [
"title_lengths:0", "title_encoder:0", "ratings:0", "query_lengths:0",
"passage_lengths:0", "features:0", "encoder:0", "decoder:0",
"Placeholder:0"
]
for name in input_names:
inputs_schema.append({IOSchemaLoader.NAME_KEY: name})
outputs_schema = [{
IOSchemaLoader.NAME_KEY: "output_identity:0"
}, {
IOSchemaLoader.NAME_KEY: "loss_identity:0"
}]
onnx_model_path = f"_{str(uuid.uuid4())}.onnx"
# generate input data
data_path = os.path.join(TF_MODEL_PATH, 'test_data.npz')
save_test_data_to_disk(data_path)
cvt_config = ConversionConfig(model_path=os.path.join(
os.path.dirname(__file__), 'data', 'full_doran_frozen.pb'),
inputs_schema=inputs_schema,
outputs_schema=outputs_schema,
sample_input_data_path=data_path,
model_framework="tensorflow",
onnx_model_path=onnx_model_path)
convert(cvt_config)
os.remove(onnx_model_path)
def test_tf_frozen_pb():
model_path = os.path.join(os.path.dirname(__file__), "models",
"full_doran_frozen.pb")
test_data_path = os.path.join(os.path.dirname(__file__), "models",
"doran.npz")
inputs_schema = [{
"name": "title_lengths:0"
}, {
"name": "title_encoder:0"
}, {
"name": "ratings:0"
}, {
"name": "query_lengths:0"
}, {
"name": "passage_lengths:0"
}, {
"name": "features:0"
}, {
"name": "encoder:0"
}, {
"name": "decoder:0"
}, {
"name": "Placeholder:0"
}]
outputs_schema = [{
"name": "output_identity:0"
}, {
"name": "loss_identity:0"
}]
cvt_config = ConversionConfig(model_path=model_path,
inputs_schema=inputs_schema,
outputs_schema=outputs_schema,
model_framework="tensorflow",
sample_input_data_path=test_data_path)
convert(cvt_config)
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_saved_model():
logging.info("converting... saved_model")
onnx_model_path = f"_{str(uuid.uuid4())}.onnx"
cvt_config = ConversionConfig(model_path=os.path.join(
TF_MODEL_PATH, "saved_model"),
model_framework="tensorflow",
onnx_model_path=onnx_model_path)
convert(cvt_config)
os.remove(onnx_model_path)
def test_pytorch():
model_path = os.path.join(os.path.dirname(__file__), "models",
"squeezenet1_1.pth")
inputs_schema = [{"name": "input_0", "shape": [1, 3, 244, 244]}]
outputs_schema = [{"name": "output_0"}]
cvt_config = ConversionConfig(model_path=model_path,
inputs_schema=inputs_schema,
outputs_schema=outputs_schema,
model_framework="pytorch")
convert(cvt_config)
def test_tf_checkpoint():
model_path = os.path.join(os.path.dirname(__file__), "models", "tf_ckpt",
"model.ckpt.meta")
inputs_schema = [{"name": "input:0"}]
outputs_schema = [{"name": "result:0"}]
cvt_config = ConversionConfig(model_path=model_path,
inputs_schema=inputs_schema,
outputs_schema=outputs_schema,
model_framework="tensorflow")
convert(cvt_config)
def test_tf_saved_model():
model_path = os.path.join(os.path.dirname(__file__), "models",
"saved_model")
inputs_schema = [{"name": "X:0"}]
outputs_schema = [{"name": "pred:0"}]
cvt_config = ConversionConfig(model_path=model_path,
inputs_schema=inputs_schema,
outputs_schema=outputs_schema,
model_framework="tensorflow")
convert(cvt_config)
def test_frozen_graph():
logging.info("converting... frozen_graph")
inputs_schema = SCHEMAS[0]
outputs_schema = SCHEMAS[1]
onnx_model_path = f"_{str(uuid.uuid4())}.onnx"
cvt_config = ConversionConfig(model_path=os.path.join(
TF_MODEL_PATH, "mnist_frozen.pb"),
inputs_schema=inputs_schema,
outputs_schema=outputs_schema,
model_framework="tensorflow",
onnx_model_path=onnx_model_path)
convert(cvt_config)
os.remove(onnx_model_path)
def test_customized_model_with_dynamic_input(use_sample_data,
use_dynamic_input_schema):
model_dir = os.path.join(TEST_INPUT_DIR,
"test_customized_model_with_dynamic_input")
model_path = os.path.join(model_dir, "sample_model.pth")
if os.path.exists(model_dir):
shutil.rmtree(model_dir)
# prepare model and sample input data
shutil.copytree(
os.path.join(os.path.dirname(__file__), 'data', 'sample',
'sample_model'), model_dir)
sample_input_data_path = os.path.join(model_dir, 'data.npz')
sm.save_sample_input_data(sample_input_data_path)
# prepare model
inputs_schema = [{
IOSchemaLoader.NAME_KEY:
sm.INPUT_NAME,
IOSchemaLoader.SHAPE_KEY:
sm.get_input_shape(use_dynamic_input_schema),
IOSchemaLoader.DTYPE_KEY:
"float"
}]
outputs_schema = [{
IOSchemaLoader.NAME_KEY:
sm.OUTPUT_NAME,
IOSchemaLoader.SHAPE_KEY:
sm.get_output_shape(use_dynamic_input_schema)
}]
onnx_model_path = f"_{str(uuid.uuid4())}.onnx"
if use_sample_data:
cvt_config = ConversionConfig(
model_path=model_path,
inputs_schema=inputs_schema,
outputs_schema=outputs_schema,
sample_input_data_path=sample_input_data_path,
onnx_opset=11,
model_framework="pytorch",
onnx_model_path=onnx_model_path)
else:
cvt_config = ConversionConfig(model_path=model_path,
inputs_schema=inputs_schema,
outputs_schema=outputs_schema,
onnx_opset=11,
model_framework="pytorch",
onnx_model_path=onnx_model_path)
convert(cvt_config)
os.remove(onnx_model_path)
def test_checkpoint():
# test the checkpoint model with .meta file
logging.info("converting... checkpoint")
inputs_schema = SCHEMAS[0]
outputs_schema = SCHEMAS[1]
onnx_model_path = f"_{str(uuid.uuid4())}.onnx"
cvt_config = ConversionConfig(model_path=os.path.join(
TF_MODEL_PATH, "ckpt", 'model.ckpt.meta'),
inputs_schema=inputs_schema,
outputs_schema=outputs_schema,
model_framework="tensorflow",
onnx_model_path=onnx_model_path)
convert(cvt_config)
os.remove(onnx_model_path)
def test_saved_model_with_npz():
logging.info("converting... saved_model")
inputs_schema = SCHEMAS[0]
outputs_schema = SCHEMAS[1]
onnx_model_path = f"_{str(uuid.uuid4())}.onnx"
cvt_config = ConversionConfig(model_path=os.path.join(
TF_MODEL_PATH, "saved_model"),
inputs_schema=inputs_schema,
outputs_schema=outputs_schema,
sample_input_data_path=TF_INPUT_DATA,
model_framework="tensorflow",
onnx_model_path=onnx_model_path)
convert(cvt_config)
os.remove(onnx_model_path)
def test_frozen_graph_with_negative_shape():
logging.info("converting... frozen_graph")
inputs_schema = [{
IOSchemaLoader.NAME_KEY: "input:0",
IOSchemaLoader.SHAPE_KEY: [-1, 784]
}]
outputs_schema = SCHEMAS[1]
onnx_model_path = f"_{str(uuid.uuid4())}.onnx"
cvt_config = ConversionConfig(model_path=os.path.join(
TF_MODEL_PATH, "mnist_frozen.pb"),
inputs_schema=inputs_schema,
outputs_schema=outputs_schema,
model_framework="tensorflow",
onnx_model_path=onnx_model_path)
convert(cvt_config)
os.remove(onnx_model_path)
def test_infogan_pb():
logging.info("converting... infogan model")
# convert
inputs_schema = [{
IOSchemaLoader.NAME_KEY: "zc_vectors:0",
IOSchemaLoader.DTYPE_KEY: "float32",
IOSchemaLoader.SHAPE_KEY: [10, 74]
}, {
IOSchemaLoader.NAME_KEY: "is_training_generator:0",
IOSchemaLoader.DTYPE_KEY: "bool",
IOSchemaLoader.SHAPE_KEY: []
}]
outputs_schema = [{IOSchemaLoader.NAME_KEY: "generator/layer_3/Sigmoid:0"}]
test_data_path = os.path.join(os.path.dirname(__file__), "data",
"infogan_inputs.npz")
data = np.load(test_data_path, allow_pickle=True)
for x in np.ndenumerate(data):
real_data1 = np.array(x[1]['zc_vectors:0'], dtype=np.float32)
real_data2 = np.array(x[1]['is_training_generator:0'])
data = [real_data1, real_data2]
break
name_list = ["zc_vectors:0", "is_training_generator:0"]
test_data = dict(zip(name_list, data))
npz_file_path = os.path.join(os.path.dirname(__file__), "data",
"test_infogan_inputs.npz")
np.savez(npz_file_path, **test_data)
onnx_model_path = f"{str(uuid.uuid4())}.onnx"
cvt_config = ConversionConfig(model_path=os.path.join(
os.path.dirname(__file__), 'data', "infogan_after_tf_optimize.pb"),
inputs_schema=inputs_schema,
outputs_schema=outputs_schema,
sample_input_data_path=npz_file_path,
model_framework="tensorflow",
onnx_model_path=onnx_model_path)
convert(cvt_config)
# make sure the first input of the converted ONNX model is dynamic
import onnxruntime as ort
sess = ort.InferenceSession(onnx_model_path,
providers=["CPUExecutionProvider"])
assert 'unk' in sess.get_inputs()[0].shape[0]
assert 74 == sess.get_inputs()[0].shape[1]
os.remove(onnx_model_path)
def test_customized_model_with_sub_dir():
model_dir = os.path.join(TEST_INPUT_DIR,
"test_customized_model_with_sub_dir")
if os.path.exists(model_dir):
shutil.rmtree(model_dir)
# prepare model
shutil.copytree(
os.path.join(os.path.dirname(__file__), 'data', 'sample',
"multi_file_model"), os.path.join(model_dir))
model_path = os.path.join(model_dir, "data/multi_file_model.pth")
# prepare model
inputs_schema = []
for name, shape in zip(mf.get_input_names(), mf.get_input_shapes()):
inputs_schema.append({
IOSchemaLoader.NAME_KEY: name,
IOSchemaLoader.SHAPE_KEY: shape
})
outputs_schema = []
for name in mf.get_output_names():
outputs_schema.append({IOSchemaLoader.NAME_KEY: name})
onnx_model_path = f"_{str(uuid.uuid4())}.onnx"
cvt_config = ConversionConfig(model_path=model_path,
inputs_schema=inputs_schema,
outputs_schema=outputs_schema,
model_root_path=os.path.join(
model_dir, 'src'),
model_framework="pytorch",
onnx_opset=11,
onnx_model_path=onnx_model_path)
convert(cvt_config)
os.remove(onnx_model_path)
# exclude script if it's already appended during conversion
for p in sys.path:
if p.startswith(model_dir):
sys.path.remove(p)
def test_script_module():
# prepare model
model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
'models', 'scriptmodule.pt')
inputs_schema = [{
IOSchemaLoader.SHAPE_KEY: [1],
IOSchemaLoader.NAME_KEY: "input_0"
}]
outputs_schema = [{
IOSchemaLoader.SHAPE_KEY: [1],
IOSchemaLoader.NAME_KEY: "output_0"
}]
onnx_model_path = f"_{str(uuid.uuid4())}.onnx"
# test
cvt_config = ConversionConfig(model_path=model_path,
inputs_schema=inputs_schema,
outputs_schema=outputs_schema,
onnx_opset=11,
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)