def conv2d(a,
weight,
bias=None,
padding='valid',
stride=(1, 1),
dilation=(1, 1),
groups=1):
"""
NCHW
:param a: NCHW
:param weight: OIHW
:param bias:
:param padding:
:param stride: sH, sW
:param dilation: dH, dW
:param groups:
:return:
"""
assert padding in ('valid', 'same', 'full') or isinstance(
padding, (tuple, list))
if isinstance(stride, int):
stride = (stride, stride)
assert isinstance(stride, (tuple, list))
if isinstance(dilation, int):
dilation = (dilation, dilation)
assert isinstance(dilation, (tuple, list))
bs, xch, xh, xw = a.shape
zch, _, k0, k1 = weight.shape
if isinstance(padding, (tuple, list)):
a = Tensor._op(Pad, a, op_args=(padding, 'zeros'))
if padding is 'same':
zshape = np.ceil([xh / stride[0], xw / stride[1]]).astype(int)
if stride[0] < k0:
ph = (k0 - 1) * dilation[0] + zshape[0] * stride[0] - xh
else:
ph = zshape[0] * stride[0] - xh
if stride[1] < k1:
pw = (k1 - 1) * dilation[1] + zshape[1] * stride[1] - xw
else:
pw = zshape[1] * stride[1] - xw
# padding = (pw//2, (pw+1)//2, ph//2, (ph+1)//2)
padding = (ph // 2, (ph + 1) // 2, pw // 2, (pw + 1) // 2)
# print(padding, pw, ph, zshape, xh, xw, stride)
a = Tensor._op(Pad, a, op_args=(padding, 'zeros'))
out = Tensor._op(Conv2d, a, weight, op_args=(stride, dilation))
if bias is not None:
out = out + bias
return out
def max_pool2d(a, pool_size=(2, 2), stride=(2, 2), padding='valid'):
"""
:param a:
:param pool_size:
:param stride:
:param padding:
:return:
"""
assert padding in ('valid', 'same', 'full') or isinstance(
padding, (tuple, list))
if isinstance(stride, int):
stride = (stride, stride)
assert isinstance(stride, (tuple, list))
if isinstance(pool_size, int):
pool_size = (pool_size, pool_size)
assert isinstance(pool_size, (tuple, list))
bs, xch, xh, xw = a.shape
if isinstance(padding, (tuple, list)):
a = Tensor._op(Pad, a, op_args=(padding, 'zeros'))
# if padding is 'same':
# zshape = np.ceil([xh / stride[0], xw / stride[1]]).astype(int)
# if stride[0] < pool_size[0]:
# ph = (pool_size[0]-1) * dilation[0] + zshape[0] * stride[0] - xh
# else:
# ph = zshape[0] * stride[0] - xh
# if stride[1] < pool_size[1]:
# pw = (pool_size[1]-1) * dilation[1] + zshape[1] * stride[1] - xw
# else:
# pw = zshape[1] * stride[1] - xw
# # padding = (pw//2, (pw+1)//2, ph//2, (ph+1)//2)
# padding = (ph//2, (ph+1)//2, pw//2, (pw+1)//2)
# print(padding, pw, ph, zshape, xh, xw, stride)
# a = Tensor._op(Pad, a, op_args=(padding, 'zeros'))
out = Tensor._op(MaxPool2d, a, op_args=(
pool_size,
stride,
))
# return out
return out
def pad(a, expanded_padding, mode='zeros'):
return Tensor._op(Pad, a, op_args=(expanded_padding, mode))
def softmax(a, dim=None):
if dim is None:
dim = -1
return Tensor._op(Softmax, a, op_args=(dim, ))
def sigmoid(a):
return Tensor._op(Sigmoid, a)
def relu(a, inplace=False):
return Tensor._op(ReLU, a)
