def build(x):
return [
mb.avg_pool(
x=x,
kernel_sizes=kernel_sizes[:num_dims],
strides=strides[:num_dims],
pad_type="valid",
),
mb.avg_pool(
x=x,
kernel_sizes=kernel_sizes[-num_dims:],
strides=strides[-num_dims:],
pad_type="same",
exclude_padding_from_average=True,
),
]
def prog(x):
x1 = mb.transpose(x=x, perm=[0, 2, 3, 1])
x2 = mb.relu(x=x1)
out1 = mb.avg_pool(x=x2,
kernel_sizes=[1, 1],
strides=[1, 1],
pad_type="valid")
x3 = mb.transpose(x=x2, perm=[0, 3, 1, 2])
out2 = mb.log(x=x3)
return out1, out2
