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)