conv_args not in llops interface
MODIFY
tinygrad/llops/ops_cpu.py
tinygrad/llops/ops_cpu.py
import numpy as np
-from tinygrad.helpers import UnaryOps, BinaryOps, ReduceOps
+from tinygrad.helpers import get_conv_args, UnaryOps, BinaryOps, ReduceOps
class CPUBuffer(np.ndarray):
def relu(x): return np.maximum(x, 0)
@@ -77,7 +77,8 @@ def get_tx(x, C):
writeable=False,
)
-def conv(x,w,ret,C):
+def conv(x,w,ret,stride,groups):
+ C = get_conv_args(x.shape, w.shape, stride, groups)
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)
@@ -87,7 +88,8 @@ def conv(x,w,ret,C):
ret[:] = np.moveaxis(tmp,4,2).reshape(C.bs, C.groups*C.rcout, C.oy, C.ox)
return ret
-def convdw(x,grad_output,dw,C):
+def convdw(x,grad_output,dw,stride,groups):
+ C = get_conv_args(x.shape, dw.shape, stride, groups)
tx = get_tx(x, C)
ggg = grad_output.reshape(C.bs, C.groups, C.rcout, C.oy, C.ox)
gdw = dw.reshape((C.groups, C.rcout, C.cin, C.H, C.W))
@@ -97,7 +99,8 @@ def convdw(x,grad_output,dw,C):
gdw[g] += np.tensordot(ggg[:,g], tx[:,g], ((0,2,3),(0,2,3)))
return dw
-def convdx(w,grad_output,dx,C):
+def convdx(w,grad_output,dx,stride,groups):
+ C = get_conv_args(x.shape, dw.shape, stride, groups)
ggg = grad_output.reshape(C.bs, C.groups, C.rcout, C.oy, C.ox)
tw = w.reshape(C.groups, C.rcout, C.cin, C.H, C.W)
gdx = dx.reshape((C.bs, C.groups, C.cin, C.iy, C.ix))

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
import functools
import numpy as np
import pyopencl as cl
-from tinygrad.helpers import binary_broadcast, UnaryOps, BinaryOps, ReduceOps
+from tinygrad.helpers import binary_broadcast, get_conv_args, UnaryOps, BinaryOps, ReduceOps
cl_ctx, cl_queue = None, None
def require_init_gpu():
@@ -228,7 +228,8 @@ def matmul(a, b, c, transpose_a=False, transpose_b=False):
osize)
return c
-def conv(x,w,ret,C):
+def conv(x,w,ret,stride,groups):
+ C = get_conv_args(x.shape, w.shape, stride, groups)
# input = (bs, groups, cin, iy, ix)
# weight = (groups, rcout, cin, H, W)
# output = (bs, groups, rcout, oy, ox)
@@ -262,7 +263,8 @@ def conv(x,w,ret,C):
# tensw = (groups*rcout, cin, H, W)
# ggg = (bs, groups*rout, oy, ox)
-def convdw(x,grad_output,dw,C):
+def convdw(x,grad_output,dw,stride,groups):
+ C = get_conv_args(x.shape, dw.shape, stride, groups)
convdw_prg = clbuild("convdw", """
__kernel void convdw(__global const float *tensx, __global const float *ggg, __global float *dw,
int H, int W, int groups, int rcout, int cin, int oy, int ox, int iy, int ix, int ys, int xs, int bs) {
@@ -285,7 +287,8 @@ def convdw(x,grad_output,dw,C):
convdw_prg([C.groups*C.rcout*C.cin, C.H, C.W], None, x.cl, grad_output.cl, dw.cl, *[i32(x) for x in C])
return dw
-def convdx(w,grad_output,dx,C):
+def convdx(w,grad_output,dx,stride,groups):
+ C = get_conv_args(dx.shape, w.shape, stride, groups)
convdx_prg = clbuild("convdx", """
__kernel void convdx(__global const float *tensw, __global const float *ggg, __global float *dx,
int H, int W, int groups, int rcout, int cin, int oy, int ox, int iy, int ix, int ys, int xs, int bs) {

MODIFY
tinygrad/llops/ops_torch.py
tinygrad/llops/ops_torch.py
@@ -23,14 +23,14 @@ from tinygrad.llops.ops_cpu import unary_op, binary_op, reduce_op, reshape, perm
# ************* processing ops *************
-def conv(x,w,ret,C):
- ret[:] = torch.nn.functional.conv2d(x, w, stride=(C.ys,C.xs), groups=C.groups)
+def conv(x,w,ret,stride,groups):
+ ret[:] = torch.nn.functional.conv2d(x, w, stride=stride, groups=groups)
return ret
-def convdw(x,grad_output,dw,C):
- dw[:] = torch.nn.grad.conv2d_weight(x, dw.shape, grad_output, stride=(C.ys,C.xs), groups=C.groups)
+def convdw(x,grad_output,dw,stride,groups):
+ dw[:] = torch.nn.grad.conv2d_weight(x, dw.shape, grad_output, stride=stride, groups=groups)
return dw
-def convdx(w,grad_output,dx,C):
- dx[:] = torch.nn.grad.conv2d_input(dx.shape, w, grad_output, stride=(C.ys,C.xs), groups=C.groups)
+def convdx(w,grad_output,dx,stride,groups):
+ dx[:] = torch.nn.grad.conv2d_input(dx.shape, w, grad_output, stride=stride, groups=groups)
return dx

MODIFY
tinygrad/mlops.py
tinygrad/mlops.py
@@ -174,11 +174,11 @@ class Matmul(Function):
class Conv2D(Function):
def forward(ctx, x, w, stride=1, groups=1):
C = get_conv_args(x.shape, w.shape, stride, groups)
- ctx.save_for_backward(x,w,C)
- return ctx.op.conv(x, w, ctx.buffer((C.bs, C.groups*C.rcout, C.oy, C.ox)), C)
+ ctx.save_for_backward(x,w,(C.ys,C.xs), C.groups)
+ return ctx.op.conv(x, w, ctx.buffer((C.bs, C.groups*C.rcout, C.oy, C.ox)), (C.ys,C.xs), C.groups)
def backward(ctx, grad_output):
- x, w, C = ctx.saved_tensors
- dx = ctx.op.convdx(w, grad_output, ctx.buffer((C.bs, C.groups*C.cin, C.iy, C.ix)), C) if ctx.needs_input_grad[0] else None
- dw = ctx.op.convdw(x, grad_output, ctx.buffer((C.groups*C.rcout, C.cin, C.H, C.W)), C) if ctx.needs_input_grad[1] else None
+ x, w, stride, groups = ctx.saved_tensors
+ dx = ctx.op.convdx(w, grad_output, ctx.buffer(x.shape), stride, groups) if ctx.needs_input_grad[0] else None
+ dw = ctx.op.convdw(x, grad_output, ctx.buffer(w.shape), stride, groups) if ctx.needs_input_grad[1] else None
return dx, dw
\ No newline at end of file

conv_args not in llops interface