def test_avg_pool2d_no_count_pad(): kh, kw = (4, 4) sh, sw = (2, 2) ph, pw = (2, 2) x = sym.Variable("x") y = sym.avg_pool2d(x, pool_size=(kh, kw), strides=(sw, sw), padding=(ph, pw), name="y", count_include_pad=False) dtype = "float32" n = 1 (ic, ih, iw) = (3, 28, 28) (oc, oh, ow) = (3, 15, 15) a_np = np.random.uniform(low=0.001, size=(n, ic, ih, iw)).astype(dtype) pad_np = np.zeros(shape=(n, ic, ih+2*ph, iw+2*pw)).astype(dtype) no_zero = (range(n), range(ic), (range(ph, ih+ph)), (range(pw, iw+pw))) pad_np[np.ix_(*no_zero)] = a_np b_np = np.zeros(shape=(n, oc, oh, ow)).astype(dtype) for i in range(oh): for j in range(ow): pad_count = np.sum(pad_np[:, :, i*sh:i*sh+kh, j*sw:j*sw+kw] > 0, axis=(2,3)) b_np[:,:,i,j] = np.sum(pad_np[:, :, i*sh:i*sh+kh, j*sw:j*sw+kw], axis=(2,3)) / np.maximum(pad_count, 1) b_np = np.maximum(b_np, 0.0) shape_dict = {"x": (n, ic, ih, iw)} for target, ctx in ctx_list(): graph, lib, _ = nnvm.compiler.build(y, target, shape_dict) m = graph_runtime.create(graph, lib, ctx) data = tvm.nd.array(a_np) m.run(x=data) out = m.get_output(0, tvm.nd.empty((n, oc, oh, ow), dtype)) np.testing.assert_allclose(out.asnumpy(), b_np, rtol=1e-5)
def test_avg_pool2d_no_count_pad(): kh, kw = (4, 4) sh, sw = (2, 2) ph, pw = (2, 2) x = sym.Variable("x") y = sym.avg_pool2d(x, pool_size=(kh, kw), strides=(sw, sw), padding=(ph, pw), name="y", count_include_pad=False) dtype = "float32" n = 1 (ic, ih, iw) = (3, 28, 28) (oc, oh, ow) = (3, 15, 15) a_np = np.random.uniform(low=0.001, size=(n, ic, ih, iw)).astype(dtype) pad_np = np.zeros(shape=(n, ic, ih+2*ph, iw+2*pw)).astype(dtype) no_zero = (range(n), range(ic), (range(ph, ih+ph)), (range(pw, iw+pw))) pad_np[np.ix_(*no_zero)] = a_np b_np = np.zeros(shape=(n, oc, oh, ow)).astype(dtype) for i in range(oh): for j in range(ow): pad_count = np.sum(pad_np[:, :, i*sh:i*sh+kh, j*sw:j*sw+kw] > 0, axis=(2,3)) b_np[:,:,i,j] = np.sum(pad_np[:, :, i*sh:i*sh+kh, j*sw:j*sw+kw], axis=(2,3)) / np.maximum(pad_count, 1) b_np = np.maximum(b_np, 0.0) shape_dict = {"x": (n, ic, ih, iw)} for target, ctx in ctx_list(): graph, lib, _ = nnvm.compiler.build(y, target, shape_dict) m = graph_runtime.create(graph, lib, ctx) data = tvm.nd.array(a_np) m.run(x=data) out = m.get_output(0, tvm.nd.empty((n, oc, oh, ow), dtype)) tvm.testing.assert_allclose(out.asnumpy(), b_np, rtol=1e-5)
def _alter_avg_pool2d_layout(attrs, inputs, tinfo): import nnvm.symbol as sym copy_inputs = [s for s in inputs] new_attrs = {k: attrs[k] for k in attrs.keys()} # NHWC -> NCHW if attrs["layout"] != "NHWC": return None new_attrs["layout"] = "NCHW" if "target" in new_attrs: del new_attrs["target"] return sym.avg_pool2d(*copy_inputs, **new_attrs)
def test_avg_pool2d(): x = sym.Variable("x") y = sym.avg_pool2d(x, pool_size=(2,2), strides=(2,2), padding=(0,0), name="y") dtype = "float32" dshape = (1, 3, 28, 28) oshape = (1, 3, 14, 14) shape_dict = {"x": dshape} for target, ctx in ctx_list(): graph, lib, _ = nnvm.compiler.build(y, target, shape_dict) m = graph_runtime.create(graph, lib, ctx) data = tvm.nd.array(np.random.uniform(size=dshape).astype(dtype)) m.run(x=data) out = m.get_output(0, tvm.nd.empty(oshape, dtype)) b_np = np.mean(data.asnumpy().reshape(1,3,14,2,14,2), axis=(3,5)) np.testing.assert_allclose(out.asnumpy(), b_np, rtol=1e-5)
def test_avg_pool2d(): x = sym.Variable("x") y = sym.avg_pool2d(x, pool_size=(2,2), strides=(2,2), padding=(0,0), name="y") dtype = "float32" dshape = (1, 3, 28, 28) oshape = (1, 3, 14, 14) shape_dict = {"x": dshape} for target, ctx in ctx_list(): graph, lib, _ = nnvm.compiler.build(y, target, shape_dict) m = graph_runtime.create(graph, lib, ctx) data = tvm.nd.array(np.random.uniform(size=dshape).astype(dtype)) m.run(x=data) out = m.get_output(0, tvm.nd.empty(oshape, dtype)) b_np = np.mean(data.asnumpy().reshape(1,3,14,2,14,2), axis=(3,5)) tvm.testing.assert_allclose(out.asnumpy(), b_np, rtol=1e-5)
def Pooling(data, kernel, stride, pad, pool_type, name): if pad[0] != 0 or pad[1] != 0: data = sym.pad(data=data, pad_width=((0, 0), (pad[0], pad[0]), (pad[1], pad[1]), (0, 0))) if pool_type == 'max': return sym.max_pool2d(data=data, pool_size=kernel, strides=stride, name=name, layout='NHWC') if pool_type == 'avg': return sym.avg_pool2d(data=data, pool_size=kernel, strides=stride, name=name, layout='NHWC') raise ValueError("Invalid pooling type: " + pool_type)
import nnvm import nnvm.compiler import nnvm.symbol as sym data = sym.Variable('data') weight = sym.Variable('weight') c1 = sym.conv2d(data=data, weight=weight, use_bias=False, channels=32, kernel_size=(4, 4)) p1 = sym.avg_pool2d(data=c1, pool_size=(4, 4)) s1 = sym.relu(data=p1) compute_graph = nnvm.graph.create(s1) print('-------original') print(compute_graph.ir()) print('-------after') deploy_graph, lib, params = nnvm.compiler.build( compute_graph, target='llvm', shape={'data': (1, 32, 32, 32)}, dtype='float32') print(deploy_graph.ir())
def avg_pool2d_callback(attrs, inputs, tinfos): print("AVG_POOL2D!!!") new_attrs = {k: attrs[k] for k in attrs.keys()} new_attrs['layout'] = 'NCHWc' return sym.avg_pool2d(inputs[0], **new_attrs)