def global_max_pool2d_callback(attrs, inputs, tinfos):
print("GLOBAL_MAX_POOL2D!!!")
input = inputs[0]
# data = tinfos[0]
# if len(data.shape) == 5:
# _, C, H, W, c = [x.value for x in data.shape]
# FIXME: hack for last layer
input = sym.data_reorder_back(input)
return sym.global_max_pool2d(input)
def _alter_global_max_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.global_max_pool2d(*copy_inputs, **new_attrs)
def test_alter_func_return_none():
data = sym.Variable("data", shape=(1, 32, 512, 512))
pool1 = sym.global_max_pool2d(data, name="pool1", layout="NCHW")
pool2 = sym.global_max_pool2d(pool1, name="pool2", layout="NCHW")
relu = sym.relu(pool2, name="relu")
g = graph.create(relu)
g = g.apply("CorrectLayout")
g = graph_attr.set_dtype_inputs(g, "float32")
g = g.apply(["InferShape", "InferType"])
assert g.json_attr("layout") == ['NCHW', 'NCHW', 'NCHW', 'NCHW']
@reg.register_alter_op_layout("global_max_pool2d", level=100)
def alter_global_max_pool2d_layout(attrs, inputs, tinfos):
return None
g = g.apply("AlterOpLayout")
# alter func return none, nothing get replaced,
# the layouts should remain the same
assert g.json_attr("layout") == ['NCHW', 'NCHW', 'NCHW', 'NCHW']
def test_global_max_pool2d():
x = sym.Variable("x")
y = sym.global_max_pool2d(x, name="y")
dtype = "float32"
dshape = (1, 1024, 7, 7)
oshape = (1, 1024, 1, 1)
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.max(data.asnumpy(), axis=(2,3), keepdims=True)
np.testing.assert_allclose(out.asnumpy(), b_np, rtol=1e-5)
def test_global_max_pool2d():
x = sym.Variable("x")
y = sym.global_max_pool2d(x, name="y")
dtype = "float32"
dshape = (1, 1024, 7, 7)
oshape = (1, 1024, 1, 1)
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.max(data.asnumpy(), axis=(2,3), keepdims=True)
tvm.testing.assert_allclose(out.asnumpy(), b_np, rtol=1e-5)
def test_global_pool2d():
x = sym.Variable("data", shape=(1, 32, 512, 512))
y = sym.global_max_pool2d(x, name="pool", layout="NCHW")
g, ldict = correct_layout(y)
assert (ldict["data"][0] == "NCHW")
assert (ldict["pool"][0] == "NCHW")
# if index of H and W remain the same,
# pool2d does not convert the layout.
g, ldict = correct_layout(g, "NCHW16c")
assert (ldict["data"][0] == "NCHW16c")
assert (ldict["pool"][0] == "NCHW16c")
# for other layout it requires a layout transform.
g, ldict = correct_layout(g, "NHWC")
assert (ldict["data"][0] == "NHWC")
assert (ldict["data_NCHW"][0] == "NCHW")
assert (ldict["pool"][0] == "NCHW")
def test_global_pool2d():
x = sym.Variable("data", shape=(1, 32, 512, 512))
y = sym.global_max_pool2d(x, name="pool", layout="NCHW")
g, ldict = correct_layout(y)
assert(ldict["data"][0] == "NCHW")
assert(ldict["pool"][0] == "NCHW")
# if index of H and W remain the same,
# pool2d does not convert the layout.
g, ldict = correct_layout(g, "NCHW16c")
assert(ldict["data"][0] == "NCHW16c")
assert(ldict["pool"][0] == "NCHW16c")
# for other layout it requires a layout transform.
g, ldict = correct_layout(g, "NHWC")
assert(ldict["data"][0] == "NHWC")
assert(ldict["data_NCHW"][0] == "NCHW")
assert(ldict["pool"][0] == "NCHW")
def check(in_shape, out_shape, **kwargs):
x = sym.Variable("x", shape=in_shape)
y = sym.global_max_pool2d(x, name="y", **kwargs)
sdict = infer_shape(y)
assert(tuple(sdict["y"][0]) == tuple(out_shape))
def test_max_pool2d():
x = sym.Variable('x')
y = sym.max_pool2d(x, pool_size=(3, 3), name="y")
y = sym.global_max_pool2d(y)
assert y.list_input_names() == ["x"]
def check(in_shape, out_shape, **kwargs):
x = sym.Variable("x", shape=in_shape)
y = sym.global_max_pool2d(x, name="y", **kwargs)
sdict = infer_shape(y)
assert(tuple(sdict["y"][0]) == tuple(out_shape))
def test_max_pool2d():
x = sym.Variable('x')
y = sym.max_pool2d(x, pool_size=(3, 3), name="y")
y = sym.global_max_pool2d(y)
assert y.list_input_names() == ["x"]
