diff --git a/topi/python/topi/math.py b/topi/python/topi/math.py index 2113f6f4da3f..c24ef5163bba 100644 --- a/topi/python/topi/math.py +++ b/topi/python/topi/math.py @@ -2,6 +2,42 @@ from __future__ import absolute_import as _abs import tvm +@tvm.tag_scope(tag='ewise') +def identity(x): + """Take identity of input x. + + Parameters + ---------- + x : tvm.Tensor + Input argument. + + Returns + ------- + y : tvm.Tensor + The result. + """ + # pylint: disable=unnecessary-lambda + return tvm.compute(x.shape, lambda *i: x(*i)) + + +@tvm.tag_scope(tag='ewise') +def negation(x): + """Take negation of input x. + + Parameters + ---------- + x : tvm.Tensor + Input argument. + + Returns + ------- + y : tvm.Tensor + The result. + """ + # pylint: disable=unnecessary-lambda + return tvm.compute(x.shape, lambda *i: -x(*i)) + + @tvm.tag_scope(tag="ewise") def exp(x): """Take exponential of input x. diff --git a/topi/python/topi/nn/__init__.py b/topi/python/topi/nn/__init__.py index 50a0cf8e53ef..3c3ebb3f6d0a 100644 --- a/topi/python/topi/nn/__init__.py +++ b/topi/python/topi/nn/__init__.py @@ -2,7 +2,12 @@ """Neural network operators""" from __future__ import absolute_import as _abs -from .ewise import * -from .mapping import * -from .conv import * +from .batch_norm import * +from .convolution import * +from .elemwise import * from .dilate import * +from .flatten import * +from .fully_connected import * +from .mapping import * +from .pooling import * +from .softmax import * diff --git a/topi/python/topi/nn/batch_norm.py b/topi/python/topi/nn/batch_norm.py new file mode 100644 index 000000000000..8a75460a0d69 --- /dev/null +++ b/topi/python/topi/nn/batch_norm.py @@ -0,0 +1,55 @@ +"""TVM operator batch normalization compute.""" +from __future__ import absolute_import +import tvm + +@tvm.tag_scope(tag='batch_norm') +def batch_norm(data, gamma, beta, moving_mean, moving_var, eps, fix_gamma): + """Batch normalization operator in NCHW layout. + + Parameters + ---------- + data : tvm.Tensor + 4-D with shape [batch, channel, height, width] + + gamma : tvm.Tensor + 1-D with shape [channel] + + beta : tvm.Tensor + 1-D with shape [channel] + + moving_mean : tvm.Tensor + 1-D with shape [channel] + + moving_var : tvm.Tensor + 1-D with shape [channel] + + eps : float + Epsilon to prevent div 0. + + fix_gamma : boolean + Fix gamma while training + + Returns + ------- + output : tvm.Tensor + 4-D with shape [batch, channel, height, width] + + mean : tvm.Tensor + 1-D with shape [channel] + + var : tvm.Tensor + 1-D with shape [channel] + """ + assert len(data.shape) == 4, "only support 4-dim batch norm" + batch, channel, height, width = data.shape + if fix_gamma: + out = tvm.compute((batch, channel, height, width), \ + lambda b, c, h, w: (data[b, c, h, w] - moving_mean[c]) / \ + tvm.intrin.sqrt(moving_var[c] + eps) + beta[c]) + else: + out = tvm.compute((batch, channel, height, width), \ + lambda b, c, h, w: (data[b, c, h, w] - moving_mean[c]) / \ + tvm.intrin.sqrt(moving_var[c] + eps) * gamma[c] + beta[c]) + mean = tvm.compute((C, ), lambda c: moving_mean[c]) + var = tvm.compute((C, ), lambda c: moving_var[c]) + return [out, mean, var] diff --git a/topi/python/topi/nn/conv.py b/topi/python/topi/nn/convolution.py similarity index 99% rename from topi/python/topi/nn/conv.py rename to topi/python/topi/nn/convolution.py index 504af1aa24b5..8d82d4d8ab3a 100644 --- a/topi/python/topi/nn/conv.py +++ b/topi/python/topi/nn/convolution.py @@ -89,7 +89,7 @@ def conv2d_hwcn(Input, Filter, stride, padding): Returns ------- - Output : tvm.Tensor + output : tvm.Tensor 4-D with shape [out_height, out_width, out_channel, batch] """ assert isinstance(stride, int) or len(stride) == 2 diff --git a/topi/python/topi/nn/ewise.py b/topi/python/topi/nn/elemwise.py similarity index 100% rename from topi/python/topi/nn/ewise.py rename to topi/python/topi/nn/elemwise.py diff --git a/topi/python/topi/nn/flatten.py b/topi/python/topi/nn/flatten.py new file mode 100644 index 000000000000..e75ab73216b2 --- /dev/null +++ b/topi/python/topi/nn/flatten.py @@ -0,0 +1,32 @@ +"""TVM operator flatten compute.""" +from __future__ import absolute_import +import tvm + +@tvm.tag_scope(tag='flatten') +def flatten(data): + """Flattens the input array into a 2-D array by collapsing the higher dimensions. + + Parameters + ---------- + data : tvm.Tensor + Input array. + + Returns + ------- + output : tvm.Tensor + 2-D array with collapsed higher dimensions. + """ + ishape = data.shape + dim = 1 + for i in range(1, len(ishape)): + dim = dim * ishape[i] + oshape = [ishape[0], dim] + + def unwrap(idx, shape): + index = [] + for s in reversed(shape): + index.append(idx % s) + idx = idx / s + return list(reversed(index)) + + return tvm.compute(oshape, lambda i, j: data(i, *unwrap(j, ishape[1:]))) diff --git a/topi/python/topi/nn/fully_connected.py b/topi/python/topi/nn/fully_connected.py new file mode 100644 index 000000000000..870df02424f5 --- /dev/null +++ b/topi/python/topi/nn/fully_connected.py @@ -0,0 +1,62 @@ +"""TVM operator fully connected compute.""" +from __future__ import absolute_import +import tvm + + +@tvm.tag_scope(tag='fully_connected') +def fully_connected(data, weight): + """Matrix multiplication + + Parameters + ---------- + data : tvm.Tensor + 2-D with shape [batch, in_dim] + + weight : tvm.Tensor + 2-D with shape [out_dim, in_dim] + + Returns + ------- + output : tvm.Tensor + 2-D with shape [batch, out_dim] + """ + assert len(data.shape) == 2 and len(weight.shape) == 2, \ + "only support 2-dim fully_connected" + batch, in_dim = data.shape + out_dim, _ = weight.shape + k = tvm.reduce_axis((0, in_dim), name='k') + return tvm.compute((batch, out_dim), lambda i, j: \ + tvm.sum(data[i][k] * weight[j][k], axis=k)) + + +@tvm.tag_scope(tag='fully_connected_with_bias') +def fully_connected_with_bias(data, weight, bias): + """Applies a linear transformation: :math:`Y = XW^T + b`. + + Parameters + ---------- + data : tvm.Tensor + 2-D with shape [batch, in_dim] + + weight : tvm.Tensor + 2-D with shape [out_dim, in_dim] + + bias : tvm.Tensor + 1-D with shape [out_dim] + + Returns + ------- + output : tvm.Tensor + 2-D with shape [batch, out_dim] + """ + assert len(data.shape) == 2 and len(weight.shape) == 2, \ + "only support 2-dim fully_connected" + assert len(data.shape) == 2 and len(weight.shape) == 2 and len(bias.shape) == 1, \ + "only support 2-dim fully_connected" + batch, in_dim = data.shape + out_dim, _ = weight.shape + k = tvm.reduce_axis((0, in_dim), name='k') + matmul = tvm.compute((batch, out_dim), lambda i, j: \ + tvm.sum(data[i, k] * weight[j, k], axis=k)) + return tvm.compute((batch, out_dim), lambda i, j: \ + matmul[i, j] + bias[j]) diff --git a/topi/python/topi/nn/pooling.py b/topi/python/topi/nn/pooling.py new file mode 100644 index 000000000000..9f8322da3b21 --- /dev/null +++ b/topi/python/topi/nn/pooling.py @@ -0,0 +1,77 @@ +"""TVM operator pooling compute.""" +from __future__ import absolute_import +import tvm + +@tvm.tag_scope(tag='max_pool') +def max_pool(data, kernel, stride, pad): + """Perform max pooling on the data + + Parameters + ---------- + data : tvm.Tensor + 4-D with shape [batch, channel, in_height, in_width] + + kernel : list/tuple of two ints + Kernel size, or [kernel_height, kernel_width] + + stride : list/tuple of two ints + Stride size, or [stride_height, stride_width] + + pad : list/tuple of two ints + Pad size, or [pad_height, pad_width] + + Returns + ------- + output : tvm.Tensor + 4-D with shape [batch, channel, out_height, out_width] + """ + assert len(data.shape) == 4, "only support 4-dim pooling" + assert len(stride.shape) == 2, "only support 2-dim stride" + assert len(pad.shape) == 2, "only support 2-dim pad" + kernel_height, kernel_width = kernel + stride_height, stride_width = stride + pad_height, pad_width = pad + batch, channel, height, width = data.shape + padded_height = height + 2*pad_height + padded_width = width + 2*pad_width + out_height = (height + 2*pad_height - kernl_height) / stride_height + 1 + out_width = (width + 2*pad_width - kernel_width) / stride_width + 1 + + dheight = tvm.reduce_axis((0, kernel_height)) + dwidth = tvm.reduce_axis((0, kernel_width)) + + temp = tvm.compute((batch, channel, padded_height, padded_width), lambda i, c, h, w: \ + tvm.select( + tvm.make.Or(tvm.make.Or((h < pad_height), (h >= height + pad_height)), + tvm.make.Or((w < pad_width), (w >= width + pad_width))), + tvm.min_value('float32'), + data[i, c, h - pad_height, w - pad_width]), name='temp') + + return tvm.compute((batch, channel, out_height, out_width), lambda i, c, h, w: \ + tvm.max(temp[i, c, h*stride_height+dheight, w*stride_width+dwidth], axis=[dheight, dwidth])) + + +@tvm.tag_scope(tag='global_avg_pool') +def global_avg_pool(data): + """Perform global average pooling on the data + + Parameters + ---------- + data : tvm.Tensor + 4-D with shape [batch, channel, in_height, in_width] + + Returns + ------- + output : tvm.Tensor + 4-D with shape [batch, channel, 1, 1] + """ + assert len(data.shape) == 4, "only support 4-dim pooling" + batch, channel, height, width = data.shape + + dheight = tvm.reduce_axis((0, height)) + dwidth = tvm.reduce_axis((0, width)) + + tsum = tvm.compute((batch, channel, 1, 1), lambda n, c, h, w: \ + tvm.sum(data[n, c, dheight, dwidth], axis=[dheight, dwidth])) + return tvm.compute((batch, channel, 1, 1), lambda n, c, h, w: \ + tsum[n, c, h, w] / (height*width)) diff --git a/topi/python/topi/nn/softmax.py b/topi/python/topi/nn/softmax.py new file mode 100644 index 000000000000..56553441d6bc --- /dev/null +++ b/topi/python/topi/nn/softmax.py @@ -0,0 +1,27 @@ +"""TVM operator softmax compute.""" +from __future__ import absolute_import +import tvm + +@tvm.tag_scope(tag='softmax') +def softmax(x): + """Perform softmax activation on the data + + Parameters + ---------- + data : tvm.Tensor + 2-D input data + + Returns + ------- + output : tvm.Tensor + 2-D output with same shape + """ + assert len(x.shape) == 2, "only support 2-dim softmax" + m, n = x.shape + k = tvm.reduce_axis((0, n), name='k') + max_elem = tvm.compute((m, ), lambda i: \ + tvm.max(x[i, k]), axis=k) + expsum = tvm.compute((m, ), lambda i: \ + tvm.sum(tvm.exp(x[i, k] - max_elem[i]), axis=k)) + return tvm.compute(x.shape, lambda i, j: \ + tvm.exp(x[i, j] - max_elem[i]) / expsum[i])