def test_layer_autodiff(filename):
batch_size = 4
train_graph = pm.from_onnx(filename)
batch_pass = pm.UpdateBatchSize(batch_size, train_graph.name)
train_graph = batch_pass(train_graph)
target_layer = "batch_norm"
# for name, node in train_graph.nodes.items():
# if isinstance(node, pm.Template) and node.op_name == target_layer:
# for i in node.inputs:
# print(f"Input {i.name} - {i.shape}")
# print()
# for i in node.outputs:
# print(f"Output {i.name} - {i.shape}")
# break
#
# print(f"\nAfter\n")
# for name, node in train_graph.nodes.items():
# if isinstance(node, pm.Template) and node.op_name == target_layer:
# for i in node.inputs:
# print(f"Input {i.name} - {i.shape}")
# print()
# for i in node.outputs:
# print(f"Output {i.name} - {i.shape}")
# break
train_graph = pm.create_training_graph(train_graph)
def get_onnx_lenet(inp_info):
BENCH_DIR = Path(f"{Path(__file__).parent}/../benchmarks/onnx_files")
filename = f"lenet.onnx"
filepath = f"{BENCH_DIR}/full_dnns/{filename}"
assert Path(filepath).exists()
graph = pm.from_onnx(filepath)
tvm_code = pm.generate_tvm(graph, inp_info, "")
return tvm_code
def lenet_from_onnx():
filename = f"mnist-lenet.onnx"
filepath = f"{BENCH_DIR}/full_dnns/{filename}"
assert Path(filepath).exists()
graph = pm.from_onnx(filepath)
print_skip_nodes = ['write', 'output', 'var_index', 'input', 'index']
for k,v in graph.nodes.items():
if v.op_name not in print_skip_nodes:
print(f"{k} - {v.op_name}")
def convert_torch_model(input_var, model, model_name, optimize_model, training_mode, to_polymath,
convert_data_format=False):
f = io.BytesIO()
mode = torch.onnx.TrainingMode.TRAINING if training_mode else torch.onnx.TrainingMode.EVAL
if 'mask_rcnn' not in model_name:
torch.onnx.export(model, # model being run
input_var, # model input (or a tuple for multiple inputs)
f, # where to save the model (can be a file or file-like object)
export_params=True, # store the trained parameter weights inside the model file
do_constant_folding=True, # whether to execute constant folding for optimization
keep_initializers_as_inputs=True,
training=mode,
input_names=['input'], # the model's input names
output_names=['output'],
opset_version=12)
else:
model.eval()
# input_var = [(input_var,)]
if isinstance(input_var[0][-1], dict):
input_var = input_var[0] + ({},)
else:
input_var = input_var[0]
dynamic_axes = {"images_tensors": [0, 1, 2], "boxes": [0, 1], "labels": [0],
"scores": [0], "masks": [0, 1, 2]}
torch.onnx.export(model, # model being run
input_var, # model input (or a tuple for multiple inputs)
f, # where to save the model (can be a file or file-like object)
do_constant_folding=True, # whether to execute constant folding for optimization
# training=mode,
input_names=["images_tensors"],
output_names=["boxes", "labels", "scores", "masks"],
dynamic_axes=dynamic_axes,
opset_version=_onnx_opset_version,
verbose=False,
# export_params=True, # store the trained parameter weights inside the model file
# keep_initializers_as_inputs=True,
# operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN,
)
print(type(f.getvalue()))
model_proto = onnx.ModelProto.FromString(f.getvalue())
print_nodes(model_proto)
onnx.checker.check_model(model_proto)
add_value_info_for_constants(model_proto)
model_proto = onnx.shape_inference.infer_shapes(model_proto)
filepath = f"{CWD}/{model_name}.onnx"
if optimize_model:
model_proto, check = simplify(model_proto)
assert check
model_proto = update_node_names(model_proto)
model_proto = update_edge_names(model_proto)
with open(filepath, "wb") as f:
f.write(model_proto.SerializeToString())
if to_polymath:
graph = pm.from_onnx(filepath)
pm.pb_store(graph, f"{CWD}/full_dnns/")
def test_load_maskrcnn():
# mrcnn_path = f"{ONNX_DNNS}/mask_rcnn_vision_backbone.onnx"
mrcnn_path = f"{ONNX_DNNS}/resnet18_train.onnx"
graph = pm.from_onnx(mrcnn_path)
for name, node in graph.nodes.items():
if node.op_name in ['conv', 'conv_bias']:
print(node.inputs[1].name)
print(f"{node.inputs[1].init_value is not None}")
print(type(node.inputs[1]))
def test_resnet18_train():
filename = f"resnet18_train.onnx"
filepath = f"{BENCH_DIR}/full_dnns/{filename}"
assert Path(filepath).exists()
graph = pm.from_onnx(filepath)
full_path = f"{BENCH_DIR}/full_dnns"
pb_path = f"{full_path}/resnet18_train.srdfg"
pm.pb_store(graph, full_path)
node = pm.pb_load(pb_path, verbose=True)
assert len(node.nodes) == len(graph.nodes)
def test_convert_benchmarks(benchmark_name, feature_dict, data_func, input_keys, output_key):
feature_size = [str(v) for k,v in feature_dict.items()]
tabla_path = f"{OUTPATH}/{benchmark_name}_{'_'.join(feature_size)}_onnx_tabla.json"
ref_tabla_path = f"{OUTPATH}/{benchmark_name}_{'_'.join(feature_size)}_tabla.json"
filename = f"{benchmark_name}{'_'.join(feature_size)}.onnx"
filepath = f"{BENCH_DIR}/ml_algorithms/{filename}"
assert Path(filepath).exists()
graph = pm.from_onnx(filepath, use_filename=False)
# Apply transformations and/or generate verilog using 'transformed_graph'
int_feat_dict = {k: int(v) for k,v in feature_dict.items()}
_, ref_in_info, ref_out_info, ref_keys = data_func(**int_feat_dict)
int_feat_dict['coarse'] = True
ref_graph, in_info, out_info, ref_keys = data_func(**int_feat_dict)
translated_inputs = {input_keys[k]: v for k,v in in_info.items() if k in input_keys}
for i in output_key:
input_cpy = pickle.loads(pickle.dumps(translated_inputs))
np_res = out_info[i[0]]
onnx_res = graph(i[1], input_cpy)
np.testing.assert_allclose(np.squeeze(np_res), np.squeeze(onnx_res))
print(f"Starting tabla compilation\n\n")
tabla_ir, tabla_graph = pm.generate_tabla(graph,
feature_dict,
tabla_path,debug=False,
context_dict={}, add_kwargs=True)
ref_tabla_ir, ref_tabla_graph = pm.generate_tabla(ref_graph,
feature_dict,
ref_tabla_path,debug=False,
context_dict={}, add_kwargs=True)
ref_ocount_pass = pm.CountOpTypes(skip=['temp', 'parameter', ref_tabla_graph.name])
_ = ref_ocount_pass(ref_tabla_graph)
ocount_pass = pm.CountOpTypes(skip=['temp', 'parameter', 'output', 'write', tabla_graph.name])
_ = ocount_pass(tabla_graph)
pprint.pprint(ref_ocount_pass.op_types)
pprint.pprint(ocount_pass.op_types)
if set(ocount_pass.op_types.keys()) != set(ref_ocount_pass.op_types.keys()):
raise RuntimeError(f"Unequal amounts of operations for graphs:\n"
f"\tReference: {ref_ocount_pass.op_types.keys()}\n"
f"\tActual: {ocount_pass.op_types.keys()}")
for k,v in ocount_pass.op_types.items():
if v != ref_ocount_pass.op_types[k]:
raise RuntimeError(f"Unequal operations for key {k}:\n"
f"\tRef: {ref_ocount_pass.op_types[k]}\n"
f"\tActual: {v}\n")
assert len(ref_tabla_ir) == len(tabla_ir)
def test_resnet18_batchsize():
batch_size = 32
resnet18_path = f"{ONNX_DNNS}/resnet18.onnx"
resnet18_graph = pm.from_onnx(resnet18_path)
batch_size_pass = UpdateBatchSize(batch_size, resnet18_graph.op_name)
updated_resnet18 = batch_size_pass(resnet18_graph)
test_op_shape_pass = CollectDNNShapes()
_ = test_op_shape_pass(updated_resnet18)
ref_resnet18_path = f"{ONNX_DNNS}/resnet18_batch{batch_size}.onnx"
#
ref_resnet18_graph = pm.from_onnx(ref_resnet18_path)
ref_op_shape_pass = CollectDNNShapes()
_ = ref_op_shape_pass(ref_resnet18_graph)
ref_shapes = ref_op_shape_pass.shape_tracker
test_shapes = test_op_shape_pass.shape_tracker
assert len(list(ref_shapes.keys())) == len(list(test_shapes.keys())), f"Reference keys: {list(ref_shapes.keys())}\n" \
f"Test keys: {list(test_shapes.keys())}"
for op_name, shapes in ref_shapes.items():
for idx, s in enumerate(shapes):
assert isinstance(s, tuple) and s == test_shapes[op_name][idx]
def test_lenet():
filename = f"lenet.onnx"
full_path = f"{BENCH_DIR}/full_dnns"
filepath = f"{full_path}/{filename}"
pb_path = f"{full_path}/lenet.srdfg"
assert Path(filepath).exists()
graph = pm.from_onnx(filepath)
pm.pb_store(graph, full_path)
node = pm.pb_load(pb_path, verbose=True)
assert len(node.nodes) == len(graph.nodes)
for name, n in node.nodes.items():
if n.op_name == "conv":
print(n.kwargs.keys())
break
def run_onnx_benchmark(benchmark_name, feature_size):
filename = f"{benchmark_name}{'-'.join(feature_size)}.onnx"
filepath = f"{BENCH_DIR}/ml_algorithms/{filename}"
if Path(filepath).exists():
features = tuple([int(i) for i in feature_size])
graph = pm.from_onnx(filepath)
if benchmark_name == "svm":
create_svm(*features, graph)
elif benchmark_name == "reco":
create_reco(*features, graph)
elif benchmark_name == "logistic":
create_logistic(*features)
elif benchmark_name == "svm_wifi":
create_logistic(*features)
else:
raise RuntimeError(
f"Benchmark {filename} does not exist in {filepath}.")
def test_load_logistic(m):
benchmark_name = "logistic"
feature_dict = {'m': m}
input_keys = {"y":"y:0", "x":"x:0", "w":"W:0"}
output_key = [("w", "W:0")]
feature_size = [str(v) for k, v in feature_dict.items()]
filename = f"{benchmark_name}{'_'.join(feature_size)}.onnx"
filepath = f"{BENCH_DIR}/ml_algorithms/{filename}"
assert Path(filepath).exists()
graph = pm.from_onnx(filepath)
int_feat_dict = {k: int(v) for k, v in feature_dict.items()}
_, ref_in_info, ref_out_info, ref_keys = logistic(**int_feat_dict)
int_feat_dict['coarse'] = True
ref_graph, in_info, out_info, ref_keys = logistic(**int_feat_dict)
translated_inputs = {input_keys[k]: v for k,v in in_info.items() if k in input_keys}
for i in output_key:
input_cpy = pickle.loads(pickle.dumps(translated_inputs))
np_res = out_info[i[0]]
onnx_res = graph(i[1], input_cpy)
np.testing.assert_allclose(np.squeeze(np_res), np.squeeze(onnx_res))
def test_translate_linear_regressor(m):
out_key_map = {"y": "y:0", "x": "x:0", "w": "W:0"}
in_key_map = [("w", "W:0")]
fpath = f"{ONNX_FILE_DIR}/linear_{m}.onnx"
shape_dict = {"m": m}
graph = pm.from_onnx(fpath)
test_graph, input_info, out_info, keys = linear(m=m, coarse=True)
tinput_info = copy.deepcopy(input_info)
tkeys = copy.deepcopy(keys)
test_res = test_graph(tkeys, tinput_info)
np.testing.assert_allclose(test_res, (out_info["w"]))
onx_input_info = copy.deepcopy(input_info)
translated_inputs = {out_key_map[k]: v for k,v in input_info.items() if k in out_key_map}
onnx_res = graph(in_key_map[0][1], translated_inputs)
np.testing.assert_allclose(onnx_res, (out_info["w"]))
tabla_path = f"{OUTPATH}/{graph.name}{m}_tabla.json"
tabla_ir = pm.generate_tabla(graph,
shape_dict,
tabla_path)
def test_autodiff():
resnet18_path = f"{ONNX_DNNS}/resnet18.onnx"
resnet18_graph = pm.from_onnx(resnet18_path)
def test_load_maskrcnn():
# mrcnn_path = f"{ONNX_DNNS}/mask_rcnn_vision_backbone.onnx"
mrcnn_path = f"{ONNX_DNNS}/resnet18_train.onnx"
graph = pm.from_onnx(mrcnn_path)
def convert_model_to_polymath(model_path):
graph = pm.from_onnx(model_path)
root_path = Path(model_path).parent
pm.pb_store(graph, f"{root_path}/srdfg/")
def test_layer_autodiff(filename):
graph = pm.from_onnx(filename)
# train_graph = graph
train_graph = pm.create_training_graph(graph)
layout_pass = pm.UpdateLayout('nchw', 'nhwc')
train_graph = layout_pass(train_graph)
