def test_argmax():
ib = relay.ir_builder.IRBuilder()
n, c, h, w = tvm.var("n"), tvm.var("c"), tvm.var("h"), tvm.var("w")
x = ib.param("x", relay.ty.TensorType((n, c, h, w), "float32"))
with ib.function(x) as func:
ib.ret(relay.argmax(x, axis=(1, )))
ib.ret(func)
func = relay.ir_pass.infer_type(ib.env, func.to_func())
ftype = func.checked_type
assert ftype.ret_type == relay.ty.TensorType((n, h, w), "int32")
ib = relay.ir_builder.IRBuilder()
n, c, h, w = tvm.var("n"), tvm.var("c"), tvm.var("h"), tvm.var("w")
x = ib.param("x", relay.ty.TensorType((n, c, h, w), "float32"))
with ib.function(x) as func:
ib.ret(relay.argmax(x, axis=(2, ), keepdims=True))
ib.ret(func)
func = relay.ir_pass.infer_type(ib.env, func.to_func())
ftype = func.checked_type
assert ftype.ret_type == relay.ty.TensorType((n, c, 1, w), "int32")
ib = relay.ir_builder.IRBuilder()
n, c, h, w = tvm.var("n"), tvm.var("c"), tvm.var("h"), tvm.var("w")
x = ib.param("x", relay.ty.TensorType((n, c, h, w), "float32"))
with ib.function(x) as func:
ib.ret(relay.argmax(x, axis=(2, ), keepdims=True, exclude=True))
ib.ret(func)
func = relay.ir_pass.infer_type(ib.env, func.to_func())
ftype = func.checked_type
assert ftype.ret_type == relay.ty.TensorType((1, 1, h, 1), "int32")
def test_arg_reduce():
for op in [relay.argmax, relay.argmin]:
n, c, h, w = 10, 20, 3, 4
x = relay.var("x", relay.ty.TensorType((n, c, h, w), "float32"))
z = relay.argmax(x, axis=(1, ))
"axis=" in z.astext()
zz = relay.ir_pass.infer_type(z)
assert zz.checked_type == relay.ty.TensorType((n, h, w), "int32")
n, c, h, w = tvm.var("n"), tvm.var("c"), tvm.var("h"), tvm.var("w")
x = relay.var("x", relay.ty.TensorType((n, c, h, w), "float32"))
z = relay.argmax(x, axis=(2, ), keepdims=True)
zz = relay.ir_pass.infer_type(z)
assert zz.checked_type == relay.ty.TensorType((n, c, 1, w), "int32")
n, c, h, w = tvm.var("n"), tvm.var("c"), tvm.var("h"), tvm.var("w")
x = relay.var("x", relay.ty.TensorType((n, c, h, w), "float32"))
z = relay.argmax(x, axis=(2, ), keepdims=True, exclude=True)
zz = relay.ir_pass.infer_type(z)
assert zz.checked_type == relay.ty.TensorType((1, 1, h, 1), "int32")
def relay_argmax(c, v, dims):
"""Implementation of argmax for Relay."""
v = c.ref(v)
assert dims.is_constant(tuple)
return relay.cast(relay.argmax(v, axis=dims.value), "int64")
def get_net(batch_size, image_shape, num_classes, dtype):
height = image_shape[1]
width = image_shape[2]
data_shape = (batch_size,) + image_shape
net = relay.var("data", shape=data_shape, dtype=dtype)
net = conv_3x3(net, 3, 64, "conv1_1")
net = conv_3x3(net, 64, 64, "conv1_2")
net = relay.nn.max_pool2d(net, pool_size=(2, 2), strides=(2, 2), ceil_mode=True)
net = conv_3x3(net, 64, 128, "conv2_1")
net = conv_3x3(net, 64, 128, "conv2_2")
net = relay.nn.max_pool2d(net, pool_size=(2, 2), strides=(2, 2), ceil_mode=True)
net = conv_3x3(net, 128, 256, "conv3_1")
net = conv_3x3(net, 256, 256, "conv3_2")
net = conv_3x3(net, 256, 256, "conv3_3")
net = relay.nn.max_pool2d(net, pool_size=(2, 2), strides=(2, 2), ceil_mode=True)
net = conv_3x3(net, 256, 512, "conv4_1")
net = conv_3x3(net, 512, 512, "conv4_2")
net = conv_3x3(net, 512, 512, "conv4_3")
net = relay.nn.max_pool2d(net, pool_size=(2, 2), strides=(2, 2), ceil_mode=True)
net = conv_3x3(net, 512, 512, "conv5_1")
net = conv_3x3(net, 512, 512, "conv5_2")
net = conv_3x3(net, 512, 512, "conv5_3")
net = relay.nn.dropout(net, rate=0.5)
net = conv_3x3(net, 512, 72, "conv6", activation=False)
net = relay.transpose(net, (0, 2, 3, 1))
num_class_probs = 9 * 3
num_confidence_scores = 9
#num_box_delta = 9 * 4
pred_class_probs, pred_conf, pred_box_delta = relay.split(net,
(num_class_probs, num_class_probs + num_confidence_scores),
axis=-1)
# Probability
pred_class_probs = relay.reshape(pred_class_probs, (-1, 3))
pred_class_probs = relay.nn.softmax(pred_class_probs)
pred_class_probs = relay.reshape(pred_class_probs, (batch_size, -1, 3))
# Confidence
pred_conf = relay.sigmoid(pred_conf)
pred_conf = relay.reshape(pred_conf, (batch_size, -1, 1))
# Bbox_delta
pred_box_delta = relay.reshape(pred_box_delta, (batch_size, -1, 4))
delta_x, delta_y, delta_w, delta_h = relay.split(pred_box_delta, (1, 2, 3), axis=2)
delta_x = relay.reshape(delta_x, (batch_size, -1))
delta_y = relay.reshape(delta_y, (batch_size, -1))
delta_w = relay.reshape(delta_w, (batch_size, -1))
delta_h = relay.reshape(delta_h, (batch_size, -1))
anchor_box = set_anchors(height, width)
anchor_x = relay.Constant(tvm.nd.array(anchor_box[:, 0]))
anchor_y = relay.Constant(tvm.nd.array(anchor_box[:, 1]))
anchor_w = relay.Constant(tvm.nd.array(anchor_box[:, 2]))
anchor_h = relay.Constant(tvm.nd.array(anchor_box[:, 3]))
box_center_x = anchor_x + delta_x * anchor_w
box_center_y = anchor_y + delta_y * anchor_h
'''
box_width = anchor_w * relay.exp(delta_w)
box_height = anchor_h * relay.exp(delta_h)
'''
box_width = anchor_w + safe_exp(delta_w)
box_height = anchor_h + safe_exp(delta_h)
xmins, ymins, xmaxs, ymaxs = bbox_transform(box_center_x, box_center_y, box_width, box_height)
xmins = relay.minimum(relay.maximum(relay.const(0.0), xmins), relay.const(width - 1.0))
ymins = relay.minimum(relay.maximum(relay.const(0.0), ymins), relay.const(height - 1.0))
xmaxs = relay.maximum(relay.minimum(relay.const(width - 1.0), xmaxs), relay.const(0.0))
ymaxs = relay.maximum(relay.minimum(relay.const(height - 1.0), ymaxs), relay.const(0.0))
det_boxes = relay.stack(bbox_transform_inv(xmins, ymins, xmaxs, ymaxs), axis=-1)
probs = relay.multiply(pred_class_probs, pred_conf)
det_probs = relay.max(probs, axis=2)
det_class = relay.argmax(probs, axis=2)
out = relay.Tuple([det_boxes, det_probs, det_class])
args = relay.analysis.free_vars(out)
return relay.Function(args, out)