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)
