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 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 define_graph(self, inp, weight, bias, grad, inp_grad, weight_grad, bias_grad, optimizer, optimizer_kwargs): transA = False transB = False if grad.shape[1] != weight.shape[0]: indices = tuple([pm.index(0, s - 1) for s in weight.shape]) # weight_transposed = pm.temp(name=f"{weight.name}_transposed", shape=(weight.shape[1], weight.shape[0])) weight_transposed = pm.state(name=f"{weight.name}_transposed", shape=(weight.shape[1], weight.shape[0])) weight_transposed[indices[1], indices[0]] = weight[indices] pm.gemm_no_bias(grad, weight_transposed, inp_grad, transA=transA, transB=transB, strict_shapes=True) else: pm.gemm_no_bias(grad, weight, inp_grad, transA=transA, transB=transB, strict_shapes=True) if grad.shape[0] != inp.shape[1]: indices = tuple([pm.index(0, s - 1) for s in inp.shape]) # inp_transposed = pm.temp(name=f"{inp.name}_transposed", shape=(inp.shape[1], inp.shape[0])) inp_transposed = pm.state(name=f"{inp.name}_transposed", shape=(inp.shape[1], inp.shape[0])) inp_transposed[indices[1], indices[0]] = inp[indices] pm.gemm_no_bias(inp_transposed, grad, weight_grad, transA=transA, transB=transB, strict_shapes=True) else: pm.gemm_no_bias(inp, grad, weight_grad, transA=transA, transB=transB, strict_shapes=True) # Weight update assert weight_grad.shape == weight.shape OPTIMIZERS[optimizer](weight, weight_grad, **optimizer_kwargs) pm.reduce_sum(grad, bias_grad) OPTIMIZERS[optimizer](bias, bias_grad, **optimizer_kwargs)
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_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 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_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 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_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_conv2d_transpose_shapes(inp_shape, wgt_shape, stride, pad): groups = 1 dilation = 1 out_pad = 0 inp = np.random.randint(-15, 15, np.prod(inp_shape)).reshape(inp_shape) wgt = np.random.randint(-15, 15, np.prod(wgt_shape)).reshape(wgt_shape) torch_res = F.conv_transpose2d(torch.from_numpy(inp), torch.from_numpy(wgt), stride=stride, padding=pad) torch_res = conv2d_transpose(torch.from_numpy(inp), torch.from_numpy(wgt), stride, pad) # np.testing.assert_allclose(tres.numpy(), torch_res.numpy()) info = { 'data': inp, 'w': wgt, } N, C, H, W = inp.shape x = pm.input(name="data", shape=inp_shape) w = pm.state(name="w", shape=wgt_shape) out = pm.output(name="out") graph = pm.conv_transpose(x, w, out, stride, pad) # tres = graph("out", info) np.testing.assert_allclose(tres, torch_res.numpy())
def get_constant_of_shape(input_var, name=None, shape=None, out=None, **kwargs): if not out: out = pm.state(name=name, shape=shape) return out
def test_bnorm(): shape = (1, 16, 32, 32) grad = torch.rand(shape) x = torch.rand(shape) scale = torch.rand((shape[1], )) bias = torch.rand((shape[1], )) mean = torch.rand((shape[1], )) var = torch.rand((shape[1], )) torch_res = batchnorm2d_backward(grad, x, scale, bias) grad = grad.numpy() x = x.numpy() scale = scale.numpy() bias = bias.numpy() mean = mean.numpy() var = var.numpy() optimizer = "sgd" optimizer_kwargs = {"lr": 0.01} pm_x = pm.input(name="x", shape=shape) pm_grad = pm.input(name="grad", shape=shape) pm_scale = pm.state(name="scale", shape=scale.shape) pm_bias = pm.state(name="bias", shape=scale.shape) pm_mean = pm.state(name="mean", shape=scale.shape) pm_var = pm.state(name="var", shape=scale.shape) pm_x_grad = pm.output(name="x_grad", shape=shape) pm_scale_grad = pm.output(name="scale_grad", shape=scale.shape) pm_b_grad = pm.output(name="bias_grad", shape=bias.shape) inp_map = { 'x': x, 'grad': grad, 'scale': scale, 'bias': bias, 'mean': mean, 'var': var, } graph = pm.batchnorm_grad(pm_x, pm_scale, pm_bias, pm_mean, pm_var, pm_grad, pm_x_grad, pm_scale_grad, pm_b_grad, optimizer, optimizer_kwargs) rtol, atol = 1.3e-3, 1e-3 gout = graph("bias_grad", inp_map) np.testing.assert_allclose(gout, torch_res.numpy().reshape(gout.shape), rtol=rtol, atol=atol)
def define_graph(self, inp, weight, bias, grad, inp_grad, weight_grad, bias_grad, optimizer, optimizer_kwargs, stride=1, pad=0, dilation=1): min_sizes = [] k = len(grad.shape) - 2 for d in range(k): min_sizes.append((grad.shape[d + 2] - 1) * stride - 2 * pad + (weight.shape[-1] - 1) * dilation + 1) grad_input_padding = tuple(inp.shape[-k + d] - min_sizes[d] for d in range(k)) assert grad_input_padding[0] == grad_input_padding[1] pm.conv_transpose_bias(grad, weight, bias, inp_grad, stride=stride, pad=pad, out_pad=grad_input_padding[0]) inp_indices = tuple(pm.index(0, s - 1) for s in inp.shape) grad_indices = tuple(pm.index(0, s - 1) for s in grad.shape) weight_indices = tuple(pm.index(0, s - 1) for s in weight.shape) inp_transposed = pm.temp(name=f"transposed_{inp.name}", shape=(inp.shape[1], inp.shape[0], inp.shape[2], inp.shape[3])) grad_transposed = pm.state(name=f"transposed_{grad.name}", shape=(grad.shape[1], grad.shape[0], grad.shape[2], grad.shape[3])) wgt_grad_transposed = pm.temp(name=f"transposed_{weight.name}", shape=(weight.shape[1], weight.shape[0], weight.shape[2], weight.shape[3])) pm.tensor_transpose(inp, inp_transposed, perm=(1, 0, 2, 3)) pm.tensor_transpose(grad, grad_transposed, perm=(1, 0, 2, 3)) pm.conv(inp_transposed, grad_transposed, wgt_grad_transposed, stride=dilation, pad=pad, dilation=stride) pm.tensor_transpose(wgt_grad_transposed, weight_grad, perm=(1, 0, 2, 3)) # Weight update OPTIMIZERS[optimizer](weight, weight_grad, **optimizer_kwargs) pm.reduce_sum(grad, bias_grad) OPTIMIZERS[optimizer](bias, bias_grad, **optimizer_kwargs)
def test_flatten(): shape = (2, 3, 4, 5) a = np.random.random_sample(shape).astype(np.float32) for i in range(len(shape)): with pm.Node(name="flatten_op") as graph: x = pm.state("x", shape=shape) x_us = pm.flatten(x, axis=i, name="res") new_shape = (1, -1) if i == 0 else (np.prod(shape[0:i]).astype(int), -1) b = np.reshape(a, new_shape) pm_b = graph("res", {"x": a}) np.testing.assert_allclose(pm_b, b)
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 populate_state(self, node): if node.shape != pm.DEFAULT_SHAPES[0]: indices = list(product(*tuple([np.arange(i) for i in node.shape]))) for i in indices: if node.init_value is not None: init_val = node.init_value[i] else: init_val = None x = pm.state(graph=node, init_value=init_val, name=f"{node.name}{i}", root_name=node.name, shape=(1, )) self.stored_objects[id(x)] = x
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_matmul(in_shape, w_shape): x = np.random.randint(0, 30, np.prod(in_shape)).reshape(in_shape) w = np.random.randint(0, 30, np.prod(w_shape)).reshape(w_shape) if in_shape[-1] == w_shape[-1]: o_np = [email protected] else: assert in_shape[-1] == w_shape[0] o_np = x @ w with pm.Node(name="mmul") as graph: x_pm = pm.input(name="x", shape=in_shape) w_pm = pm.state(name="w", shape=w_shape) o_pm = pm.output(name="o", shape=o_np.shape) pm.matmul(x_pm, w_pm, o_pm) in_dict = {"x": x, "w": w} res = graph("o", in_dict) np.testing.assert_allclose(o_np, res)
def test_loss(shape): inp = np.random.uniform(-15, 15, np.prod(shape)).reshape(shape) tgt = np.random.randint(0, 15, np.prod(shape[0])) torch_res = F.cross_entropy(torch.from_numpy(inp), torch.from_numpy(tgt)) info = { 'data': inp, 'tgt': tgt, } np_res = torch_ce_loss(inp, tgt) np.testing.assert_allclose(np_res, torch_res.numpy()) x = pm.input(name="data", shape=shape) tgt_ = pm.state(name="tgt", shape=(shape[0], )) loss = pm.output(name="loss") graph = pm.cross_entropy_loss(x, tgt_, loss) tres = graph("loss", info) np.testing.assert_allclose(tres, np_res)
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 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_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 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
def test_flatten_reco(): with pm.Node(name="recommender") as graph: m = pm.parameter("m") n = pm.parameter("n") k = pm.parameter("k") x1 = pm.input("x1", shape=(k, )) x2 = pm.input("x2", shape=(k, )) r1 = pm.input("r1", shape=(m, )) y1 = pm.input("y1", shape=(m, )) r2 = pm.input("r2", shape=(n, )) y2 = pm.input("y2", shape=(n, )) w1 = pm.state("w1", shape=(m, k)) w2 = pm.state("w2", shape=(n, k)) i = pm.index(0, m - 1, name="i") j = pm.index(0, n - 1, name="j") l = pm.index(0, k - 1, name="l") h1_sum = pm.sum([l], (w1[i, l] * x2[l]).set_name("w1*x2")).set_name("h1_sum") h1 = (h1_sum[i] * r1[i]).set_name("h1") h2_sum = pm.sum([l], (w2[j, l] * x1[l]).set_name("w2*x1")).set_name("h2_sum") h2 = (h2_sum[j] * r2[j]).set_name("h2") d1 = (h1[i] - y1[i]).set_name("d1") d2 = (h2[j] - y2[j]).set_name("d2") g1 = (d1[i] * x2[l]).set_name("g1") g2 = (d2[j] * x1[l]).set_name("g2") w1[i, l] = (w1[i, l] - g1[i, l]) w2[j, l] = (w2[j, l] - g2[j, l]) m_ = 3 n_ = 3 k_ = 2 input_info = {} input_info["m"] = m_ input_info["n"] = n_ input_info["k"] = k_ input_info["w1"] = np.random.randint(1, 6, m_ * k_).reshape(m_, k_) input_info["w2"] = np.random.randint(1, 6, n_ * k_).reshape(n_, k_) input_info["x1"] = np.random.randint(1, 6, k_) input_info["x2"] = np.random.randint(1, 6, k_) input_info["r1"] = np.random.randint(0, 2, m_) input_info["y1"] = np.random.randint(0, 6, m_) input_info["r2"] = np.random.randint(0, 2, n_) input_info["y2"] = np.random.randint(0, 6, n_) out_info = numpy_reco(input_info) shape_val_pass = NormalizeGraph({"m": m_, "n": n_, "k": k_}) flatten_pass = Lower({}) new_graph = shape_val_pass(graph) test_res = new_graph(["w1", "w2"], input_info) np.testing.assert_allclose(test_res[0], out_info["w1"]) np.testing.assert_allclose(test_res[1], out_info["w2"]) flattened_g = flatten_pass(new_graph) input_info = {} input_info["m"] = m_ input_info["n"] = n_ input_info["k"] = k_ input_info["w1"] = np.random.randint(1, 6, m_ * k_).reshape(m_, k_) input_info["w2"] = np.random.randint(1, 6, n_ * k_).reshape(n_, k_) input_info["x1"] = np.random.randint(1, 6, k_) input_info["x2"] = np.random.randint(1, 6, k_) input_info["r1"] = np.random.randint(0, 2, m_) input_info["y1"] = np.random.randint(0, 6, m_) input_info["r2"] = np.random.randint(0, 2, n_) input_info["y2"] = np.random.randint(0, 6, n_) new_out_info = numpy_reco(input_info) pairs_w1 = list( product(*tuple([np.arange(i) for i in input_info["w1"].shape]))) pairs_w2 = list( product(*tuple([np.arange(i) for i in input_info["w2"].shape]))) w1_init = input_info["w1"] for p in pairs_w1: input_info[f"w1/w1({p[0]}, {p[1]})"] = input_info["w1"][p] input_info.pop("w1") w2_init = input_info["w2"] for p in pairs_w2: input_info[f"w2/w2({p[0]}, {p[1]})"] = input_info["w2"][p] input_info.pop("w2") for p in range(k_): input_info[f"x1/x1({p},)"] = input_info["x1"][p] input_info[f"x2/x2({p},)"] = input_info["x2"][p] input_info.pop("x1") input_info.pop("x2") for p in range(m_): input_info[f"r1/r1({p},)"] = input_info["r1"][p] input_info[f"y1/y1({p},)"] = input_info["y1"][p] input_info.pop("r1") input_info.pop("y1") for p in range(n_): input_info[f"r2/r2({p},)"] = input_info["r2"][p] input_info[f"y2/y2({p},)"] = input_info["y2"][p] input_info.pop("r2") input_info.pop("y2") w1_keys = [f"w1/w1({p[0]}, {p[1]})" for p in pairs_w1] w2_keys = [f"w2/w2({p[0]}, {p[1]})" for p in pairs_w2] all_vals = flattened_g(w1_keys + w2_keys, input_info) out1 = np.asarray(list(all_vals[0:6])).reshape(new_out_info["w2"].shape) out2 = np.asarray(list(all_vals[6:])).reshape(new_out_info["w2"].shape) np.testing.assert_allclose( new_out_info["w1"], np.asarray(list(all_vals[0:6])).reshape(new_out_info["w2"].shape)) np.testing.assert_allclose( new_out_info["w2"], np.asarray(list(all_vals[6:])).reshape(new_out_info["w2"].shape))
def generate_srdfg(onnx_graph): names = [des.name for des in onnx_graph.DESCRIPTOR.fields] graph_name = getattr(onnx_graph, "name") initializers = get_initializers(onnx_graph.initializer) mgdfg = pm.Node(name=graph_name) # TODO: This is a hotfix for identifying gradient updates, but weights should have initializers state_variables = get_states_by_gradient(onnx_graph) node_info = {} # TODO: If a value has an initializer, set the initializer value as the value for the node for o in onnx_graph.output: assert o.name not in node_info if o.name in state_variables: node_info[o.name] = pm.state(name=state_variables[o.name], shape=get_value_info_shape(o, mgdfg), graph=mgdfg) node_info[state_variables[o.name]] = node_info[o.name] else: node_info[o.name] = pm.output(name=o.name, shape=get_value_info_shape(o, mgdfg), graph=mgdfg) for i in onnx_graph.input: if i.name in state_variables.values(): assert i.name in node_info continue assert i.name not in node_info if i.name in state_variables: node_info[i.name] = pm.state(name=state_variables[i.name], shape=get_value_info_shape(i, mgdfg), graph=mgdfg) node_info[state_variables[i.name]] = node_info[i.name] elif i.name in initializers and not itercheck(initializers[i.name]): node_info[i.name] = pm.parameter(name=i.name, default=initializers[i.name], graph=mgdfg) elif i.name in initializers: node_info[i.name] = pm.state(name=i.name, shape=get_value_info_shape(i, mgdfg), graph=mgdfg) else: node_info[i.name] = pm.input(name=i.name, shape=get_value_info_shape(i, mgdfg), graph=mgdfg) for v in onnx_graph.value_info: if v.name in node_info: continue elif v.name in initializers: node_info[v.name] = pm.variable(initializers[v.name], name=v.name, shape=get_value_info_shape( v, mgdfg), graph=mgdfg) else: node_info[v.name] = { "name": v.name, "shape": get_value_info_shape(v, mgdfg) } for k, v in initializers.items(): if k not in node_info: # TODO: Need to set the value here node_info[k] = pm.state(name=k, shape=get_value_info_shape(v, mgdfg), graph=mgdfg) state_variables[k] = node_info[k] for k, v in mgdfg.nodes.items(): if isinstance(v, pm.parameter) and k not in node_info: node_info[k] = v for n in onnx_graph.node: assert n.op_type in NODE_NAMES _ = convert_node(n, mgdfg, node_info, state_variables) return mgdfg