def test_linear_deserialize():
graph_name = "linear_reg1"
with pm.Node(name=graph_name) as graph:
m = pm.placeholder("m")
x_ = pm.placeholder("x", shape=(m))
y_ = pm.placeholder("y")
w_ = pm.placeholder("w", shape=(m))
mu = pm.parameter(name="mu", default=1.0)
i = pm.index(0, (m - 1).set_name("m-1"), name="i")
h = pm.sum([i], (x_[i] * w_[i]).set_name("x*w"), name="h")
d = (h - y_).set_name("h-y")
g = (d * x_[i]).set_name("d*x")
mug = (mu * g[i]).set_name("mu*g[i]")
w_ = ((w_[i]) - mug).set_name("w_out")
x = np.random.randint(0, 10, 10)
y = np.random.randint(0, 10, 1)[0]
w = np.random.randint(0, 10, 10)
graph_res = graph("w_out", {"x": x, "y": y, "w": w})
actual_res = w - ((np.sum(x * w) - y) * x) * 1.0
np.testing.assert_allclose(graph_res, actual_res)
cwd = Path(f"{__file__}").parent
base_path = f"{cwd}/pmlang_examples"
full_path = f"{base_path}/outputs"
pb_path = f"{full_path}/{graph_name}.srdfg"
pm.pb_store(graph, full_path)
node = pm.pb_load(pb_path)
new_graph_res = node("w_out", {"x": x, "y": y, "w": w})
np.testing.assert_allclose(graph_res, new_graph_res)
np.testing.assert_allclose(actual_res, new_graph_res)
def test_single_dim_norm():
with pm.Node(name="elem1") as graph:
m = pm.parameter("m")
x = pm.input("x", shape=m)
w = pm.state("w", shape=m)
i = pm.index(0, m - 1, name="i")
w[i] = (w[i] * x[i])
x_ = np.random.randint(0, 10, 3)
w_ = np.random.randint(0, 10, 3)
coarse_eval = graph("w", x=x_, w=w_)
np_result = x_ * w_
np.testing.assert_allclose(coarse_eval, np_result)
shape_pass = NormalizeGraph({"m": 3})
graph_shapes = shape_pass(graph)
shape_res = graph_shapes("w", x=x_, w=w_)
np.testing.assert_allclose(shape_res, np_result)
lower_pass = Lower({})
lowered_graph = lower_pass(graph_shapes)
input_info = {f"w/w({i},)": w_[i] for i in range(len(w_))}
input_info.update({f"x/x({i},)": x_[i] for i in range(len(x_))})
fine_grained_eval = lowered_graph("w/w(1,)", input_info)
assert fine_grained_eval == np_result[1]
pb_path = f"{OUTPATH}/{graph.name}.srdfg"
pm.pb_store(lowered_graph, OUTPATH)
loaded_node = pm.pb_load(pb_path)
input_info = {f"w/w({i},)": w_[i] for i in range(len(w_))}
input_info.update({f"x/x({i},)": x_[i] for i in range(len(x_))})
fine_grained_eval = loaded_node("w/w(1,)", input_info)
assert fine_grained_eval == np_result[1]
def test_linear_reg():
m_ = 3
graph, input_info, out_info, keys = linear(m=m_, coarse=True)
coarse_eval = graph(keys, input_info)
np.testing.assert_allclose(coarse_eval, out_info["w"])
fgraph, input_info, out_info, keys = linear(m=m_, coarse=False)
lower_pass = Lower({})
lowered_graph = lower_pass(fgraph, {})
all_vals = lowered_graph(keys, input_info)
out = np.asarray(all_vals).reshape(out_info["w"].shape)
np.testing.assert_allclose(out, out_info["w"])
cwd = Path(f"{__file__}").parent
base_path = f"{cwd}/pmlang_examples"
full_path = f"{base_path}/outputs"
pb_path = f"{full_path}/{graph.name}.srdfg"
pm.pb_store(lowered_graph, full_path)
loaded_node = pm.pb_load(pb_path)
_, input_info, out_info, keys = linear(m=m_, coarse=False)
loaded_res = loaded_node(keys, input_info)
out = np.asarray(loaded_res).reshape(out_info["w"].shape)
np.testing.assert_allclose(out, out_info["w"])
def test_conv_embedded_values(x_shape, w_shape, params):
shape_dict = {
"n": x_shape[0],
"ic": x_shape[1],
"ih": x_shape[2],
"iw": x_shape[3],
"nf": w_shape[0],
"kh": w_shape[2],
"kw": w_shape[3],
"stride": params["stride"],
"pad": params["pad"]
}
graph, input_info0, out_info, keys = conv(x_shape,
w_shape,
params,
coarse=True,
debug_matrix=True)
ngraph, input_info1, out_info, keys = conv(x_shape,
w_shape,
params,
coarse=False,
debug_matrix=True)
lower_pass = pm.Lower({})
lowered = lower_pass(ngraph)
pb_path = f"{OUTPATH}/{graph.name}.srdfg"
pm.pb_store(lowered, OUTPATH)
node = pm.pb_load(pb_path)
assert len(node.nodes) == len(lowered.nodes)
assert list(node.nodes.keys()) == list(lowered.nodes.keys())
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_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_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 test_svm_wifi(lr, delta, features, locations, train_size):
shape_dict = {"n_locations": locations, "n_features": features}
graph, input_info0, out_info, keys = svm_wifi(features, locations, coarse=True)
tabla_path = f"{OUTPATH}/{graph.name}_{features}_{locations}_tabla.json"
# res0 = graph(keys, input_info0)[0]
# np.testing.assert_allclose(res0, out_info['weights'])
ngraph, input_info1, out_info, keys = svm_wifi(features, locations, coarse=False)
# tabla_path = f"{OUTPATH}/{graph.name}_{locations}_{features}_tabla.json"
tabla_ir, tabla_graph = pm.generate_tabla(graph,
shape_dict,
tabla_path,
context_dict=input_info1, add_kwargs=True)
srdfg_path = f"{OUTPATH}/"
pm.pb_store(tabla_graph, srdfg_path)
def test_tabla_linear(m_):
shape_dict = {"m": m_}
graph, input_info, out_info, keys = linear(m=m_, coarse=True)
lgraph, input_info, out_info, keys = linear(m=m_, coarse=False)
cwd = Path(f"{__file__}").parent
base_path = f"{cwd}/pmlang_examples"
full_path = f"{base_path}/outputs"
graph_name = f"{graph.name}_{m_}"
tabla_path = f"{full_path}/{graph_name}_tabla.json"
tabla_ir, tabla_graph = pm.generate_tabla(graph,
shape_dict,
tabla_path,
context_dict=input_info,
add_kwargs=True)
cwd = Path(f"{__file__}").parent
base_path = f"{cwd}/pmlang_examples"
full_path = f"{base_path}/outputs"
pb_path = f"{full_path}/{graph.name}.srdfg"
pm.pb_store(graph, full_path)
node = pm.pb_load(pb_path)
def test_lower_group_op():
with pm.Node(name="linear_reg1") as graph:
m = pm.parameter(name="m")
x = pm.input("x", shape=(m))
y = pm.input("y")
w = pm.state("w", shape=(m))
i = pm.index(0, m - 1, name="i")
h = pm.sum([i], w[i] * x[i], name="h")
m_ = 3
n_ = 3
x_ = np.random.randint(0, 10, m_)
w_ = np.random.randint(0, 10, (m_))
np_result = np.sum(x_ * w_)
np.testing.assert_allclose(graph("h", {"w": w_, "x": x_}), np_result)
np.testing.assert_allclose(graph("h", w=w_, x=x_), np_result)
shape_pass = NormalizeGraph({"m": m_, "n": n_})
graph_shapes = shape_pass(graph)
shape_res = graph_shapes("h", x=x_, w=w_)
np.testing.assert_allclose(shape_res, np_result)
lower_pass = Lower({})
lowered_graph = lower_pass(graph_shapes)
input_info = {f"w/w({i},)": w_[i] for i in range(len(w_))}
input_info.update({f"x/x({i},)": x_[i] for i in range(len(x_))})
#
fine_grained_eval = lowered_graph("h/h(4,)", input_info)
assert fine_grained_eval == np_result
pb_path = f"{OUTPATH}/linear_reg1.srdfg"
pm.pb_store(lowered_graph, OUTPATH)
loaded_node = pm.pb_load(pb_path) #
input_info = {f"w/w({i},)": w_[i] for i in range(len(w_))}
input_info.update({f"x/x({i},)": x_[i] for i in range(len(x_))})
loaded_res = loaded_node("h/h(4,)", input_info)
assert loaded_node.func_hash() == lowered_graph.func_hash()
assert loaded_res == np_result
def test_reco():
m_ = 3
n_ = 3
k_ = 2
shape_dict = {"m": n_, "k": k_, "n": n_}
graph, input_info, out_info, keys = reco(coarse=True, **shape_dict)
coarse_eval = graph(keys, input_info)
np.testing.assert_allclose(coarse_eval[0], out_info["w1"])
np.testing.assert_allclose(coarse_eval[1], out_info["w2"])
fgraph, input_info, out_info, keys = reco(coarse=False, **shape_dict)
lower_pass = Lower({})
lowered_graph = lower_pass(fgraph, {})
all_vals = lowered_graph(keys, input_info)
w1_elems = np.prod(out_info["w1"].shape)
w2_elems = np.prod(out_info["w2"].shape)
out1 = np.asarray(list(all_vals[0:w1_elems])).reshape(out_info["w1"].shape)
out2 = np.asarray(list(all_vals[w1_elems:])).reshape(out_info["w2"].shape)
np.testing.assert_allclose(out1, out_info["w1"])
np.testing.assert_allclose(out2, out_info["w2"])
cwd = Path(f"{__file__}").parent
base_path = f"{cwd}/pmlang_examples"
full_path = f"{base_path}/outputs"
pb_path = f"{full_path}/{graph.name}.srdfg"
pm.pb_store(lowered_graph, full_path)
loaded_node = pm.pb_load(pb_path)
_, input_info, out_info, keys = reco(coarse=False, **shape_dict)
loaded_res = loaded_node(keys, input_info)
lres1 = np.asarray(list(loaded_res[0:w1_elems])).reshape(out_info["w1"].shape)
lres2 = np.asarray(list(loaded_res[w1_elems:])).reshape(out_info["w2"].shape)
np.testing.assert_allclose(lres1, out_info["w1"])
np.testing.assert_allclose(lres2, out_info["w2"])
def test_svm_wifi_inference(lr, delta, features, locations, train_size):
shape_dict = {"n_locations": locations, "n_features": features}
graph, input_info0, out_info, keys = svm_wifi_inf(features, locations, coarse=True)
# tabla_path = f"{OUTPATH}/{graph.name}_{features}_{locations}_tabla.json"
#
res0 = graph(keys, input_info0)[0]
np.testing.assert_allclose(res0, out_info['scores'])
#
ngraph, input_info1, out_info, keys = svm_wifi_inf(features, locations, coarse=False)
#
# lower_pass = pm.Lower({})
# lowered = lower_pass(ngraph)
# res1 = np.asarray(lowered(keys, input_info1)).reshape(out_info["scores"].shape)
# np.testing.assert_allclose(res1, out_info["scores"])
tabla_path = f"{OUTPATH}/{graph.name}_{locations}_{features}_tabla.json"
tabla_ir, tabla_graph = pm.generate_tabla(graph,
shape_dict,
tabla_path,
context_dict=input_info1, add_kwargs=True)
srdfg_path = f"{OUTPATH}/"
pm.pb_store(tabla_graph, srdfg_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/")
