def test_multidim_sigmoid(m_):
with pm.Node(name="logistic") as graph:
m = pm.parameter(name="m")
n = pm.parameter(name="n")
x = pm.input("x", shape=(m))
w = pm.state("w", shape=(m))
i = pm.index(0, m - 1, name="i")
o = pm.sigmoid(w[i] * x[i], name="out")
x_ = np.random.randint(0, 10, m_).astype(np.float)
w_ = np.random.randint(0, 10, m_).astype(np.float)
shape_dict = {"m": m_}
input_dict = {"x": x_, "w": w_}
np_res = sigmoid((x_ * w_))
coarse_eval = graph("out", input_dict)
np.testing.assert_allclose(np_res, coarse_eval)
lowered = set_shape_and_lower(graph, shape_dict)
keys = [f"out/out({i},)" for i in range(m_)]
x_ = np.random.randint(0, 10, m_).astype(np.float)
w_ = np.random.randint(0, 10, m_).astype(np.float)
input_dict = {}
for i in range(m_):
input_dict[f"x/x({i},)"] = x_[i]
input_dict[f"w/w({i},)"] = w_[i]
np_res = sigmoid((x_ * w_))
lower_res = np.asarray(lowered(keys, input_dict)).reshape(np_res.shape)
np.testing.assert_allclose(lower_res, np_res)
def get_value_info_shape(vi, mgdfg):
if isinstance(vi, np.ndarray):
ret = vi.shape
else:
ret = []
for i, dim in enumerate(vi.type.tensor_type.shape.dim):
if hasattr(dim, 'dim_param') and dim.dim_param:
if dim.dim_param in mgdfg.nodes:
shape_node = mgdfg.nodes[dim.dim_param]
else:
shape_node = pm.parameter(name=dim.dim_param, graph=mgdfg)
d_val = shape_node
elif not dim.dim_value:
shape_node = pm.parameter(name=f"{vi.name}_dim_{i}",
graph=mgdfg)
d_val = shape_node
elif dim.dim_value > 0:
d_val = dim.dim_value
else:
continue
ret.append(d_val)
ret = tuple(ret)
return ret if len(ret) > 0 else (1, )
def test_multi_dim():
with pm.Node(name="elem4") as graph:
m = pm.parameter(name="m")
n = pm.parameter(name="n")
x = pm.input("x", shape=(m, n))
w = pm.state("w", shape=(m, n))
i = pm.index(0, m - 1, name="i")
j = pm.index(0, n - 1, name="j")
w[i, j] = (w[i, j] * x[i, j])
m_ = 3
n_ = 4
x_ = np.random.randint(0, 10, m_ * n_).reshape((m_, n_))
w_ = np.random.randint(0, 10, m_ * n_).reshape((m_, n_))
coarse_eval = graph("w", x=x_, w=w_)
np_result = x_ * w_
np.testing.assert_allclose(coarse_eval, np_result)
shape_pass = NormalizeGraph({"m": m_, "n": n_})
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 = {}
for i in range(m_):
for j in range(n_):
input_info[f"w/w({i}, {j})"] = w_[i, j]
input_info[f"x/x({i}, {j})"] = x_[i, j]
fine_grained_eval = lowered_graph("w/w(2, 3)", input_info)
assert fine_grained_eval == np_result[2, 3]
def test_load_linear_regressor(m_):
shape_dict = {"m": m_}
m = pm.parameter("m")
mu = pm.parameter(name="mu", default=1.0)
x = pm.input("x", shape=(m))
y = pm.input("y")
w = pm.state("w", shape=(m))
graph = pm.linear_regressor_train(x, w, y, mu, m)
test_graph, input_info, out_info, keys = linear(m=m_, coarse=True)
assert len(test_graph.nodes.keys()) == len(graph.nodes.keys())
assert op_counts(test_graph) == op_counts(graph)
shape_val_pass = pm.NormalizeGraph(shape_dict)
new_graph = shape_val_pass(graph)
test_res = new_graph(keys, input_info)
np.testing.assert_allclose(test_res, out_info["w"])
test_graph_lowered, input_info, new_out_info, keys = linear(m=m_)
flatten_pass = pm.Lower({})
test_flatten_pass = pm.Lower({})
flattened_g = flatten_pass(new_graph)
ref_lowered = test_flatten_pass(test_graph_lowered, {})
assert len(ref_lowered.nodes.keys()) == len(flattened_g.nodes.keys())
assert op_counts(ref_lowered) == op_counts(flattened_g)
all_vals = flattened_g(keys, input_info)
np.testing.assert_allclose(new_out_info["w"], all_vals)
def test_multi_dim_op_slice():
with pm.Node(name="elem2") as graph:
m = pm.parameter(name="m")
n = pm.parameter(name="n")
mu = pm.parameter(name="mu", default=2.0)
x = pm.input(name="x", shape=(m, n))
w = pm.state(name="w", shape=(m, n))
i = pm.index(0, m - 1, name="i")
j = pm.index(0, n - 1, name="j")
out = (x[i, j] * w[i, j]).set_name("w_out")
w[i, j] = (mu * (out[i, j] - w[i, j]))
m_ = 3
n_ = 2
x_ = np.random.randint(0, 10, m_ * n_).reshape((m_, n_))
w_ = np.random.randint(0, 10, m_ * n_).reshape((m_, n_))
coarse_eval = graph("w", x=x_, w=w_)
np_result = (x_ * w_ - w_) * 2.0
np.testing.assert_allclose(coarse_eval, np_result)
shape_pass = NormalizeGraph({"m": m_, "n": n_})
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 = {}
for i in range(m_):
for j in range(n_):
input_info[f"w/w({i}, {j})"] = w_[i, j]
input_info[f"x/x({i}, {j})"] = x_[i, j]
fine_grained_eval = lowered_graph("w/w(2, 1)", input_info)
assert fine_grained_eval == np_result[2, 1]
def test_translate_conv(x_shape, w_shape, params):
shape_dict = {"n": x_shape[0], "c": 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"]}
_, input_info, out_info, keys = conv(x_shape, w_shape, params, coarse=True, debug_matrix=False)
n = pm.parameter(name="n")
c = pm.parameter(name="ic")
ih = pm.parameter(name="ih")
iw = pm.parameter(name="iw")
nf = pm.parameter(name="nf")
kh = pm.parameter(name="kh")
kw = pm.parameter(name="kw")
x = pm.input(name="data", shape=(n, c, ih, iw))
w = pm.state(name="w", shape=(nf, c, kh, kw))
b = pm.state(name="bias", shape=(nf))
stride = pm.parameter(name="stride")
pad = pm.parameter(name="pad")
out = pm.output(name="out")
graph = pm.conv_bias(x, w, b, out, stride, pad)
tinput_info = copy.deepcopy(input_info)
res0 = graph("out", tinput_info)
np.testing.assert_allclose(res0, out_info["out"])
def test_name_change():
with pm.Node() as graph:
operation = pm.parameter(default=None, name='operation1')
pm.parameter(default=None, name='operation3')
assert 'operation1' in graph.nodes
operation.name = 'operation2'
assert 'operation2' in graph.nodes
assert graph['operation2'] is operation
# We cannot rename to an existing operation
with pytest.raises(ValueError):
operation.name = 'operation3'
def test_squeeze():
with pm.Node(name="indexop") as graph:
m = pm.parameter(name="m")
n = pm.parameter(name="n")
x = pm.state("x", shape=(m, n))
x_us = pm.squeeze(x, axis=None, name="res")
m_ = 5
n_ = 1
x_ = np.random.randint(0, 10, (m_, n_))
input_info = {"m": m_, "n": n_, "x": x_}
res = graph("res", input_info)
np.testing.assert_allclose(res, np.squeeze(x_, axis=1))
def test_conditional_callback():
with pm.Node() as graph:
a = pm.parameter(default=1)
b = pm.parameter(default=2)
c = pm.placeholder()
d = pm.predicate(c, a, b + 1)
# Check that we have "traced" the correct number of operation evaluations
tracer = pm.Profiler()
assert graph(d, {c: True}, callback=tracer) == 1
assert len(tracer.times) == 3
tracer = pm.Profiler()
assert graph(d, {c: False}, callback=tracer) == 3
assert len(tracer.times) == 4
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_dict():
expected = 13
with pm.Node() as graph:
a = pm.parameter(default=expected)
b = pm.identity({'foo': a})
actual = graph(b)['foo']
assert actual is expected, "expected %s but got %s" % (expected, actual)
def test_list():
expected = 13
with pm.Node() as graph:
a = pm.parameter(default=expected)
b = pm.identity([a, a])
actual, _ = graph(b)
assert actual is expected, "expected %s but got %s" % (expected, actual)
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_single_dim_op_slice():
with pm.Node(name="elem3") as graph:
m = pm.parameter(name="m")
x = pm.input("x", shape=m)
w = pm.state("w", shape=m)
i = pm.index(0, m - 1, name="i")
out = (w[i] * x[i])
w[i] = (out[i] - w[i])
m_ = 3
x_ = np.random.randint(0, 10, m_)
w_ = np.random.randint(0, 10, m_)
coarse_eval = graph("w", x=x_, w=w_)
np_result = x_ * w_ - 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(2,)", input_info)
assert fine_grained_eval == np_result[2]
def create_svm_wifi(features, locations, lr=0.0001, deltav=1, train_size=7703):
with pm.Node(name="svm_wifi") as graph:
learning_rate = pm.parameter("learning_rate", default=lr)
delta = pm.parameter("delta", default=deltav)
n_features = pm.parameter("n_features", default=features)
n_locations = pm.parameter("n_locations", default=locations)
x_train = pm.input("x_train", shape=(n_features, ))
y_train = pm.input("y_train", shape=(n_locations, ))
y_train_inv = pm.input("y_train_inv", shape=(n_locations, ))
weights = pm.state("weights", shape=(n_features, n_locations))
i = pm.index(0, n_features - 1, name="i")
j = pm.index(0, n_locations - 1, name="j")
scores = pm.sum([i], (weights[i, j] * x_train[i]), name="scores")
correct_class_score = pm.sum([j], (scores[j] * y_train[j]),
name="correct_class_score")
h = ((scores[j] - correct_class_score + delta).set_name("h") > 0)
# margin = (pm.cast(np.float32, h[j]) * y_train_inv[j]).set_name("margin")
margin = (h[j] * y_train_inv[j]).set_name("margin")
valid_margin_count = pm.sum([j], margin[j], name="valid_margin_count")
partial = (y_train[j] * valid_margin_count).set_name("partial")
updated_margin = (margin[j] - partial[j]).set_name("updated_margin")
# # #
dW = (x_train[i] * updated_margin[j]).set_name("dW")
weights[i, j] = (weights[i, j] -
learning_rate * dW[i, j]).set_name("weights_update")
shape_dict = {"n_features": features, "n_locations": locations}
input_info, keys, out_info = svm_wifi_datagen(features,
locations,
lr,
deltav,
lowered=True)
cwd = Path(f"{__file__}").parent
full_path = f"{cwd}/outputs"
tabla_path = f"{full_path}/{graph.name}_{locations}_{features}_tabla.json"
tabla_ir, tabla_graph = pm.generate_tabla(graph,
shape_dict,
tabla_path,
context_dict=input_info,
add_kwargs=True)
def test_sigmoid(m_):
with pm.Node(name="logistic1") as graph:
m = pm.parameter(name="m")
n = pm.parameter(name="n")
x = pm.input("x", shape=(m))
w = pm.state("w", shape=(m))
i = pm.index(0, m - 1, name="i")
o = pm.sigmoid(pm.sum([i], w[i] * x[i]), name="out")
x_ = np.random.randint(0, 10, m_)
w_ = np.random.randint(0, 10, m_)
input_dict = {"x": x_, "w": w_}
np_res = int(sigmoid(np.sum(x_ * w_)))
shape_dict = {"m": m_}
coarse_eval = graph("out", x=x_, w=w_)
np.testing.assert_allclose(np_res, coarse_eval)
lowered = set_shape_and_lower(graph, shape_dict)
def test_index_op():
with pm.Node(name="indexop") as graph:
m = pm.parameter(name="m")
n = pm.parameter(name="n")
i = pm.index(0, m - 1, name="i")
j = pm.index(0, n - 1, name="j")
i_ = (i + 1).set_name("i_")
k = (i + j).set_name("k")
m_ = 5
n_ = 3
input_info = {"m": m_, "n": n_}
res = graph("k", input_info)
op1 = np.arange(0, m_)
op2 = np.arange(0, n_)
value = np.array(list(product(*(op1, op2))))
value = np.array(list(map(lambda x: x[0] + x[1], value)))
np.testing.assert_allclose(res, value)
def test_binary_operators_right(binary_operators):
operator, a, b, expected = binary_operators
with pm.Node() as graph:
_b = pm.parameter(default=b)
operation = eval('a %s _b' % operator)
actual = graph(operation)
assert actual == expected, "expected %s %s %s == %s but got %s" % \
(a, operator, b, expected, actual)
def test_conditional():
with pm.Node() as graph:
x = pm.parameter(default=4)
y = pm.placeholder(name='y')
condition = pm.placeholder(name='condition')
z = pm.predicate(condition, x, y)
assert graph(z, condition=False, y=5) == 5
# We expect a value error if we evaluate the other branch without a placeholder
with pytest.raises(ValueError):
print(graph(z, condition=False))
def test_new():
test_a = np.array([1, 2, 3, 4])
test_b = np.array([5, 6, 7, 8])
test_placeholder = pm.placeholder("hello")
with pm.Node(name="main") as graph:
a = pm.parameter(default=6, name="a")
b = pm.parameter(default=5, name="b")
a = (a + b).set_name("a_mul_b")
with pm.Node(name="graph2") as graph2:
c = pm.variable([[[0, 1, 2], [3, 4, 5], [6, 7, 8]],
[[9, 10, 11], [12, 13, 14], [15, 16, 17]],
[[18, 19, 20], [21, 22, 23], [24, 25, 26]]],
name="c")
c_2 = (c * 2).set_name(name="c2")
e = pm.parameter(default=4, name="e")
l = pm.placeholder("test")
x = (l * e).set_name("placeholdermult")
i = pm.index(0, 1, name="i")
j = pm.index(0, 1, name="j")
k = pm.index(0, 2, name="k")
e_i = pm.var_index(c, [i, j, k], "e_i")
def linear_reg():
with pm.Node(name="linear_reg") as graph:
m = pm.placeholder("m")
x = pm.placeholder("x", shape=(m), type_modifier="input")
y = pm.placeholder("y", type_modifier="input")
w = pm.placeholder("w", shape=(m), type_modifier="state")
mu = pm.parameter(name="mu", default=1.0)
i = pm.index(0, (graph["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).set_name("d*x")
w_ = (w - (mu*g).set_name("mu*g")).set_name("w-mu*g")
def test_second():
test_a = np.array([1, 2, 3, 4])
test_b = np.array([5, 6, 7, 8])
with pm.Node(name="main") as graph:
a = pm.parameter(default=6, name="a")
b = pm.parameter(default=5, name="b")
a = (a + b).set_name("a_mul_b")
with pm.Node(name="graph2") as graph2:
n = pm.placeholder("n")
b = pm.placeholder("b")
e = pm.parameter(default=6, name="e")
l = pm.state("test", shape=(n, b))
i = pm.index(0, graph2["n"] - 1)
j = pm.index(0, graph2["b"] - 1)
lij = pm.var_index(l, [i, j], "lij")
x = (l * e).set_name("placeholdermult")
_ = graph2("test", {l: np.arange(16).reshape((-1, 4))})
_ = graph2("lij", {l: np.arange(16).reshape((-1, 4))})
_ = graph2("placeholdermult", {l: np.arange(16).reshape((-1, 4))})
def gen_from_shape(graph_type, input_shape, params=None):
if graph_type == "linear":
x = pm.input(name="x", shape=input_shape)
w = pm.state(name="w", shape=input_shape)
y = pm.input(name="y")
mu = pm.parameter(name="mu", default=1.0)
m = pm.parameter(name="m", default=input_shape)
return pm.linear_regressor_train(x,
w,
y,
mu,
m,
name="linear_regressor")
elif graph_type == "logistic":
x = pm.input(name="x", shape=input_shape)
w = pm.state(name="w", shape=input_shape)
y = pm.input(name="y")
mu = pm.parameter(name="mu", default=1.0)
m = pm.parameter(name="m", default=input_shape)
return pm.logistic_regressor_train(x,
w,
y,
mu,
m,
name="logistic_regressor")
elif graph_type == "svm":
x = pm.input(name="x", shape=input_shape)
w = pm.state(name="w", shape=input_shape)
y = pm.input(name="y")
mu = pm.parameter(name="mu", default=1.0)
m = pm.parameter(name="m", default=input_shape)
return pm.svm_classifier_train(x, w, y, mu, m, name="svm_classifier")
def test_load_nested_linear_regressor(m_):
shape_dict = {"m": m_}
with pm.Node(name="nested_linear") as graph:
m = pm.parameter(name="m")
mu = pm.parameter(name="mu", default=1.0)
x = pm.input("x", shape=(m))
y = pm.input("y")
w = pm.state("w", shape=(m))
pm.linear_regressor_train(x, w, y, mu, m, name="linear_regressor")
j = pm.index(0, m-1, name="j")
tw = (w[j] - 4).set_name("tw")
test_graph, input_info, out_info, keys = linear(m=m_, coarse=True)
shape_val_pass = pm.NormalizeGraph(shape_dict)
new_graph = shape_val_pass(graph)
test_res = new_graph("tw", input_info)
np.testing.assert_allclose(test_res, (out_info["w"] - 4))
ref_graph, input_info, new_out_info, keys = linear(m=m_)
flatten_pass = pm.Lower({})
keys = [f"tw/tw({i},)" for i in range(m_)]
flattened_g = flatten_pass(new_graph)
all_vals = flattened_g(keys, input_info)
def linear_reg_graph():
graph_name = "linear_reg"
with pm.Node(name="linear_reg") as graph:
m = pm.placeholder("m")
mu = pm.parameter(name="mu", default=1.0)
x_ = pm.placeholder("x", shape=(m), type_modifier="input")
y_ = pm.placeholder("y", type_modifier="input")
w_ = pm.placeholder("w", shape=(m), type_modifier="state")
i = pm.index(0, 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")
w_out = (w_[i]) - mu * g[i]
w_out.set_name("res")
return graph
def test_conditional_with_length():
def f(a):
return a, a
with pm.Node() as graph:
x = pm.parameter(default=4)
y = pm.placeholder(name='y')
condition = pm.placeholder(name='condition')
z1, z2 = pm.predicate(condition,
pm.func_op(f, x).set_name("xfunc"),
pm.func_op(f, y).set_name("yfunc"),
shape=2,
name="predfunc")
assert graph([z1, z2], condition=True) == (4, 4)
assert graph([z1, z2], condition=False, y=5) == (5, 5)
def test_linear_embedded():
with pm.Node(name="linear_reg") as graph:
m = pm.placeholder("m")
x = pm.placeholder("x", shape=(m), type_modifier="input")
y = pm.placeholder("y", type_modifier="input")
w = pm.placeholder("w", shape=(m), type_modifier="state")
mu = pm.parameter(name="mu", default=1.0)
i = pm.index(0, (graph["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).set_name("d*x")
w_ = (w - (mu * g).set_name("mu*g")).set_name("w-mu*g")
x = np.random.randint(0, 10, 5)
y = np.random.randint(0, 10, 1)[0]
w = np.random.randint(0, 10, 5)
graph_res = graph("w-mu*g", {"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)
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_avg_pool(data_shape, kernel_shape, stride):
data = np.random.randint(0, 5, data_shape)
tout = pooling(data, kernel_shape[0], kernel_shape[1], stride=stride)
out = pm.output(name="out")
n = pm.parameter("ns")
ic = pm.parameter("ic")
ih = pm.parameter("ih")
iw = pm.parameter("iw")
kh = pm.parameter("kh")
kw = pm.parameter("kw")
x = pm.input(name="data", shape=(n, ic, ih, iw))
g = pm.avg_pool2d(x, out, kh, kw, stride=stride, pad=0)
inp_info = {}
inp_info["data"] = data
inp_info["kh"] = kernel_shape[0]
inp_info["kw"] = kernel_shape[1]
test_out = g("out", inp_info)
np.testing.assert_allclose(test_out, tout)
def lenet(lenet_type="lenet5", coarse=True, debug=False):
with pm.Node(name="lenet") as graph:
n = pm.parameter(name="n")
c = pm.parameter(name="ic")
ih = pm.parameter(name="ih")
iw = pm.parameter(name="iw")
nf1 = pm.parameter(name="nf1")
kh1 = pm.parameter(name="kh1")
kw1 = pm.parameter(name="kw1")
data = pm.input(name="data", shape=(n, c, ih, iw))
w1 = pm.state(name="w1", shape=(nf1, c, kh1, kw1))
b1 = pm.state(name="b1", shape=(nf1))
s1 = pm.parameter(name="s1")
p1 = pm.parameter(name="p1")
c1 = pm.output(name="c1", shape=(n, nf1, 28, 28))
a1 = pm.output(name="a1", shape=(n, nf1, 28, 28))
l1 = pm.output(name="l1", shape=(n, nf1, 14, 14))
pm.conv_bias(data, w1, b1, c1, s1, p1)
pm.elem_tanh(c1, a1, shape=a1.shape)
pm.avg_pool2d(a1, l1, 2, 2, 2, 0)
nf2 = pm.parameter(name="nf2")
kh2 = pm.parameter(name="kh2")
kw2 = pm.parameter(name="kw2")
w2 = pm.state(name="w2", shape=(nf2, nf1, kh2, kw2))
b2 = pm.state(name="b2", shape=(nf2))
s2 = pm.parameter(name="s2")
p2 = pm.parameter(name="p2")
c2 = pm.output(name="c2", shape=(n, nf2, 10, 10))
a2 = pm.output(name="a2", shape=(n, nf2, 10, 10))
l2 = pm.output(name="l2", shape=(n, nf2, 5, 5))
pm.conv_bias(l1, w2, b2, c2, s2, p2)
pm.elem_tanh(c2, a2, shape=a2.shape)
pm.avg_pool2d(a2, l2, 2, 2, 2, 0)
nf3 = pm.parameter(name="nf3")
kh3 = pm.parameter(name="kh3")
kw3 = pm.parameter(name="kw3")
w3 = pm.state(name="w3", shape=(nf3, nf2, kh3, kw3))
b3 = pm.state(name="b3", shape=(nf3))
s3 = pm.parameter(name="s3")
p3 = pm.parameter(name="p3")
c3 = pm.output(name="c3", shape=(n, nf3, 1, 1))
a3 = pm.output(name="a3", shape=(n, nf3, 1, 1))
pm.conv_bias(l2, w3, b3, c3, s3, p3)
pm.elem_tanh(c3, a3, shape=a3.shape)
f4 = pm.output(name="f4", shape=(n, nf3))
pm.coarse_flatten(a3, f4, axis=1, shape=f4.shape)
m5 = pm.parameter(name="m5")
n5 = pm.parameter(name="n5")
f5 = pm.output(name="f5", shape=(n, m5))
w5 = pm.state(name="w5", shape=(m5, n5))
# w5 = pm.state(name="w5", shape=(n5, m5))
a6 = pm.output(name="a5", shape=(n, m5))
b5 = pm.state(name="b5", shape=(n5, ))
pm.gemm(f4,
w5,
b5,
f5,
shape=f5.shape,
alpha=1.0,
beta=0.0,
transA=False,
transB=True)
pm.elem_tanh(f5, a6, shape=a6.shape)
m7 = pm.parameter(name="m7")
n7 = pm.parameter(name="n7")
f7 = pm.output(name="f7", shape=(n, n7))
w7 = pm.state(name="w7", shape=(m7, n7))
# w7 = pm.state(name="w7", shape=(n7, m7))
b7 = pm.state(name="b7", shape=(n7, ))
pm.gemm(a6,
w7,
b7,
f7,
shape=f7.shape,
alpha=1.0,
beta=0.0,
transA=False,
transB=False)
out = pm.output(name="sm")
pm.softmax(f7, out, axis=1)
if coarse:
in_info, keys, out_info = np_lenetv2()
return graph, in_info, out_info, keys
else:
shape_dict = {
"n": 1,
"ic": 1,
"ih": 32,
"iw": 32,
"nf1": 6,
"kh1": 5,
"kw1": 5,
"s1": 1,
"p1": 0,
"nf2": 16,
"kh2": 5,
"kw2": 5,
"s2": 1,
"p2": 0,
"nf3": 120,
"kh3": 5,
"kw3": 5,
"s3": 1,
"p3": 0,
"m5": 120,
"n5": 84,
"m7": 84,
"n7": 10
}
shape_val_pass = pm.NormalizeGraph(shape_dict, debug=debug)
new_graph = shape_val_pass(graph)
in_info, keys, out_info = np_lenetv2(lowered=True)
return new_graph, in_info, out_info, keys
