support padding
MODIFY
test/test_ops.py
test/test_ops.py
@@ -259,14 +259,11 @@ class TestOps(unittest.TestCase):
bs = 4
cin = 3
H,W = 3,3
- with self.subTest(padding := 2):
+ for p in [2, (2,1)]:
+ with self.subTest(padding := p):
helper_test_op([(bs,cin,11,28), (4,cin,H,W)],
lambda x,w: torch.nn.functional.conv2d(x,w,padding=padding).relu(),
- lambda x,w: Tensor.conv2d(x,w,padding=padding).relu(), atol=1e-4)
- with self.subTest(padding := (2,1)):
- helper_test_op([(bs,cin,11,28), (4,cin,H,W)],
- lambda x,w: torch.nn.functional.conv2d(x,w,padding=padding).relu(),
- lambda x,w: Tensor.conv2d(x,w,padding=padding).relu(), atol=1e-4)
+ lambda x,w: Tensor.conv2d(x,w,padding=padding).relu(), atol=1e-4, forward_only=True)
def test_dilated_conv2d_forward(self):
bs = 4

MODIFY
tinygrad/llops/ops_cpu.py
tinygrad/llops/ops_cpu.py
@@ -69,6 +69,7 @@ def get_tx(x, C):
)
def conv(x,w,ret,C):
+ if C.px > 0 or C.py > 0: x = np.pad(x, [(0,0), (0,0), (C.py, C.py), (C.px, C.px)])
tx = get_tx(x, C)
tw = w.reshape(C.groups, C.rcout, C.cin, C.H, C.W)
tmp = np.zeros((C.bs,C.groups,C.oy,C.ox,C.rcout),dtype=x.dtype)

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -132,7 +132,7 @@ def conv(x,w,ret,C):
# output = (bs, groups, rcout, oy, ox)
conv_prg = clbuild("conv", """
__kernel void conv(__global const float *input, __global const float *weight, __global float *output,
- int H, int W, int groups, int rcout, int cin, int oy, int ox, int iy, int ix, int ys, int xs, int bs, int dx, int dy) {
+ int H, int W, int groups, int rcout, int cin, int oy, int ox, int iy, int ix, int ys, int xs, int bs, int dx, int dy, int px, int py) {
int B = get_global_id(0)/(groups*rcout); // range 0-bs
int g = (get_global_id(0)/rcout)%groups;
@@ -146,14 +146,17 @@ def conv(x,w,ret,C):
float acc = 0.0;
for (int ci = 0; ci < cin; ci++) {
for (int y = 0; y < H; y++) { for (int x = 0; x < W; x++) {
- acc += input[B*groups*cin*iy*ix + g*cin*iy*ix + ci*iy*ix + (y*dy+IY)*ix + (x*dx+IX)] * \
- weight[g*rcout*cin*H*W + c*cin*H*W + ci*H*W + y*W + x];
+ int idx_y = y*dy + IY - py;
+ int idx_x = x*dx + IX - px;
+ int valid = (idx_y >= 0 && idx_y < iy && idx_x >= 0 && idx_x < ix);
+ acc += valid ? input[B*groups*cin*iy*ix + g*cin*iy*ix + ci*iy*ix + idx_y*ix + idx_x] * \
+ weight[g*rcout*cin*H*W + c*cin*H*W + ci*H*W + y*W + x] : 0.0;
} }
}
output[B*groups*rcout*oy*ox + g*rcout*oy*ox + c*oy*ox + Y*ox + X] = acc;
}""")
- conv_prg([C.bs*C.groups*C.rcout, C.oy, C.ox], None, x.cl, w.cl, ret.cl, *[i32(x) for x in list(C[0:12])+[C.dx, C.dy]])
+ conv_prg([C.bs*C.groups*C.rcout, C.oy, C.ox], None, x.cl, w.cl, ret.cl, *[i32(x) for x in list(C[0:12])+[C.dx, C.dy, C.px, C.py]])
# tensx = (bs, groups*cin, iy, ix)
# tensw = (groups*rcout, cin, H, W)
@@ -189,4 +192,3 @@ def convdx(grad_output,w,dx,C):
def processing_op(op,a,b,ret,C):
if op == ProcessingOps.CONV: conv(a,b,ret,C)
elif op == ProcessingOps.CONVT: convdx(a,b,ret,C)
- elif op == ProcessingOps.CONVDW: convdw(a,b,ret,C)

MODIFY
tinygrad/llops/ops_torch.py
tinygrad/llops/ops_torch.py
@@ -26,14 +26,14 @@ from tinygrad.llops.ops_cpu import unary_op, binary_op, reduce_op, movement_op
from tinygrad.ops import ProcessingOps
def processing_op(op,x,w,ret,C):
- stride, groups, dilation = (C.ys, C.xs), C.groups, (C.dy, C.dx)
+ stride, groups, dilation, padding = (C.ys, C.xs), C.groups, (C.dy, C.dx), (C.py, C.px)
if op == ProcessingOps.CONV:
- ret[:] = torch.conv2d(x, w, stride=stride, groups=groups, dilation=dilation)
+ ret[:] = torch.conv2d(x, w, stride=stride, groups=groups, dilation=dilation, padding=padding)
elif op == ProcessingOps.CONVT:
if stride == (1,1) and dilation == (1,1):
# strided needs weird "undilation": https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
w = w.reshape(C.groups, C.rcout, C.cin, C.H, C.W).flip(3, 4).transpose(2, 1).reshape(C.groups*C.cin, C.rcout, C.H, C.W)
- ret[:] = torch.conv2d(x, w, padding=(C.H-1,C.W-1), groups=groups)
+ ret[:] = torch.conv2d(x, w, padding=(C.H-C.py-1,C.W-C.px-1), groups=groups)
else:
- output_padding = [ret.shape[d+2] - ((x.shape[d+2] - 1) * stride[d] + 1 + dilation[d] * (w.shape[d+2] - 1)) for d in range(2)]
- ret[:] = torch.conv_transpose2d(x, w, stride=stride, groups=groups, output_padding=output_padding, dilation=dilation)
+ output_padding = [ret.shape[d+2] - ((x.shape[d+2] - padding[d]*2 - 1) * stride[d] + 1 + dilation[d] * (w.shape[d+2] - 1)) for d in range(2)]
+ ret[:] = torch.conv_transpose2d(x, w, padding=padding, stride=stride, groups=groups, output_padding=output_padding, dilation=dilation)

MODIFY
tinygrad/mlops.py
tinygrad/mlops.py
@@ -167,8 +167,8 @@ class Slice(Function):
# ************* processing ops *************
class Conv2D(Function):
- def forward(ctx, x, w, stride=1, groups=1, dilation=1):
- C = get_conv_args(x.shape, w.shape, stride, groups, dilation=dilation)
+ def forward(ctx, x, w, stride=1, groups=1, dilation=1, padding=0):
+ C = get_conv_args(x.shape, w.shape, stride, groups, dilation=dilation, padding=padding)
ctx.save_for_backward(x,w,C)
return ctx.processing_op(ProcessingOps.CONV, x, w, (C.bs, C.cout, C.oy, C.ox), C)
@@ -182,7 +182,7 @@ class Conv2D(Function):
# this expand is slow
grad_output_dw = ctx.movement_op(MovementOps.EXPAND, grad_output_dw, (C.bs * C.groups, C.cin, C.rcout, C.oy, C.ox))
grad_output_dw = ctx.movement_op(MovementOps.RESHAPE, grad_output_dw, (C.bs * C.groups * C.cin * C.rcout, 1, C.oy, C.ox))
- Cdw = get_conv_args(xdw.shape, grad_output_dw.shape, stride=(C.dy, C.dx), dilation=(C.ys, C.xs), groups=C.bs*C.groups*C.cin)
+ Cdw = get_conv_args(xdw.shape, grad_output_dw.shape, padding=(C.py, C.px), stride=(C.dy, C.dx), dilation=(C.ys, C.xs), groups=C.bs*C.groups*C.cin)
grad_weight = ctx.processing_op(ProcessingOps.CONV, xdw, grad_output_dw, (Cdw.bs, Cdw.cout, Cdw.oy, Cdw.ox), Cdw)
grad_weight = ctx.movement_op(MovementOps.RESHAPE, grad_weight, (C.bs, C.groups, C.cin, C.rcout, Cdw.oy, Cdw.ox))
# sum across the batch dimension

MODIFY
tinygrad/shapetracker.py
tinygrad/shapetracker.py
@@ -90,6 +90,7 @@ class ShapeTracker:
strides = [s if x == y else 0 for s,(x,y) in zip(strides_for_shape(self.shape), zip(self.shape, new_shape))]
self.views.append(View(new_shape, strides))
+ # TODO: this is a special case of slice with strides, remove it
def flip(self, *axis):
assert all([isinstance(x, int) and x >= 0 and x < len(self.shape) for x in axis])
strides = [-s if i in axis else s for i,s in enumerate(strides_for_shape(self.shape))]

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -344,10 +344,8 @@ class Tensor:
def max_pool2d(self, kernel_size=(2,2)):
return self._pool2d(*kernel_size).max(axis=(3,5))
- def conv2d(self, weight, bias=None, padding=None, **kwargs):
- if isinstance(padding, int): padding = (padding, padding)
- x = self.pad2d((padding[1], padding[1], padding[0], padding[0])) if padding is not None else self
- ret = x._conv2d(weight, **kwargs)
+ def conv2d(self, weight, bias=None, **kwargs):
+ ret = self._conv2d(weight, **kwargs)
return ret if bias is None else ret.add(bias.reshape(shape=[1, -1, 1, 1]))
# ***** broadcasted binary ops *****

support padding