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_translate_elem_mul(x_shape): a = np.random.randint(-3, 3, x_shape) b = np.random.randint(-3, 3, x_shape) np_res = a * b graph = pm.Node("elem_mul") pm_a = pm.input(name="a", shape=x_shape, graph=graph) pm_b = pm.input(name="b", shape=x_shape, graph=graph) pm_o = pm.output(name="out", shape=x_shape, graph=graph) with graph: pm.elem_mul(pm_a, pm_b, pm_o) pm_res = graph("out", {"a": a, "b": b}) np.testing.assert_allclose(pm_res, np_res)
def test_gather1(): axis = 1 x = np.random.randn(5, 4, 3, 2).astype(np.float32) idx = np.array([0, 1, 3]) with pm.Node(name="gather_op") as graph: data = pm.input(name="input", shape=x.shape) indices = pm.input(name="indices", shape=idx.shape) out = pm.gather(data, indices, axis=axis, name="res") pm_y = graph("res", {"input": x, "indices": idx}) np_y = np.take(x, idx, axis=axis) np.testing.assert_allclose(np_y, pm_y)
def test_translate_vmul(x_shape): a = np.random.randint(-3, 3, x_shape) b = np.random.randint(-3, 3, x_shape) np_res = a.dot(b) with pm.Node("vmul") as pm_graph: pm_a = pm.input(name="a", shape=x_shape) pm_b = pm.input(name="b", shape=x_shape) pm_o = pm.output(name="o", shape=x_shape) pm_s = pm.output(name="out") pm.elem_mul(pm_a, pm_b, pm_o) pm.reduce_sum(pm_o, pm_s, axes=(0,), keepdims=0) pm_res = pm_graph("out", {"a": a, "b": b}) np.testing.assert_allclose(pm_res, np_res)
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_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_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_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_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 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 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_translate_softmax(x_shape, axis): x = np.random.randint(0, 5, x_shape).astype(np.float) data = pm.input("x", shape=x.shape) out = pm.output("out") g = pm.softmax(data, out, axis=1) res = g("out", {"x": x}) np_res = np_softmax(x, axis=1) np.testing.assert_allclose(np_res, res)
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 populate_input(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: x = pm.input(graph=node, name=f"{node.name}{i}", root_name=node.name, shape=(1, )) self.stored_objects[id(x)] = x
def test_flip(in_shape, axis): x = np.random.randn(*in_shape).astype(np.float32) with pm.Node(name="flip_op") as graph: data = pm.input(name="input", shape=x.shape) out = pm.flip(data, axis, name="res") np_y = np.flip(x, axis) pm_y = graph("res", {"input": x}) np.testing.assert_allclose(np_y, pm_y)
def test_transpose(in_shape): x = np.random.randint(0, 30, np.prod(in_shape)).reshape(in_shape) with pm.Node(name="tpose") as graph: x_pm = pm.input(name="x", shape=in_shape) pm.transpose(x_pm, (1,0), name="o") in_dict = {"x": x} res = graph("o", in_dict) np.testing.assert_allclose(x.T, res)
def test_reshape(in_shape, out_shape): x = np.zeros(in_shape).astype(np.float32) with pm.Node(name="reshape_op") as graph: data = pm.input(name="input", shape=x.shape) out = pm.reshape(data, out_shape, name="res") pm_y = graph("res", {"input": x}) np_y = np.reshape(x, out_shape) np.testing.assert_allclose(np_y, pm_y) assert np_y.shape == pm_y.shape
def test_translate_flatten(x_shape): x = np.random.randint(0, 5, x_shape) data = pm.input("x", shape=x.shape) out = pm.output("out") g = pm.batch_flatten(data, out) res = g("out", {"x": x}) print(res) print(x.reshape(-1)) np.testing.assert_allclose(res, x.reshape(-1))
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_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_translate_reduce_sum(x_shape): data = np.random.randint(-3, 3, x_shape) np_res = np.sum(data) graph = pm.Node("reduce") pm_data = pm.input(name="a", shape=x_shape, graph=graph) out = pm.output(name="out", graph=graph) axis = (0,) keepdims = 0 with graph: pm.reduce_sum(pm_data, out, axes=axis, keepdims=keepdims) pm_res = graph("out", {"a": data}) np.testing.assert_allclose(pm_res, np_res)
def test_broadcast(a_shape, b_shape, c_shape): from einops import repeat with pm.Node(name="broadcast") as graph: a = pm.input("a", shape=a_shape) b = pm.input("b", shape=b_shape) c = pm.output("c", shape=c_shape) a_idx, b_idx, c_idx = _get_elem_indices(a, b, c) c[c_idx] = a[a_idx] + b[b_idx] a_np = np.random.randint(0, 32, np.prod(a_shape)).reshape(a_shape) b_np = np.random.randint(0, 32, np.prod(b_shape)).reshape(b_shape) if len(c_shape) > 2: c_np_out = np.zeros(c_shape) else: c_np_out = np.zeros((c_shape[0], 1, c_shape[1])) a_np_t = repeat(a_np, 'i k -> i k j', j=b_shape[1]) b_np_t = repeat(b_np, 'i k -> j i k', j=a_shape[0]) actual_res = (a_np_t + b_np_t).squeeze() graph_res = graph("c", {"a": a_np, "b": b_np}) np.testing.assert_allclose(graph_res, actual_res)
def test_log_softmax(shape): inp = np.random.uniform(-15, 15, np.prod(shape)).reshape(shape) torch_res = F.log_softmax(torch.from_numpy(inp)) info = { 'data': inp, } np_res = log_softmax(inp) np.testing.assert_allclose(np_res, torch_res.numpy()) x = pm.input(name="data", shape=shape) lsmx = pm.output(name="lsmx") graph = pm.log_softmax(x, lsmx, axis=1) tres = graph("lsmx", info) np.testing.assert_allclose(tres, torch_res.numpy())
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 test_pad(in_shape, pad_start, pad_end): x = np.random.randn(*in_shape).astype(np.float32) with pm.Node(name="pad_op") as graph: data = pm.input(name="input", shape=x.shape) out = pm.pad(data, pad_start, pad_end=pad_end, name="res") if pad_end is None: padding_val = tuple((pad_start[i], pad_start[i]) for i in range(len(pad_start))) else: padding_val = tuple((pad_start[i], pad_end[i]) for i in range(len(pad_start))) np_y = np.pad(x, padding_val) pm_y = graph("res", {"input": x}) assert np_y.shape == pm_y.shape np.testing.assert_allclose(np_y, pm_y)
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_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)