simpler conv2d
MODIFY
tinygrad/ops.py
tinygrad/ops.py
@@ -162,49 +162,35 @@ class Conv2D(Function):
assert cout % ctx.groups == 0
rcout = cout//ctx.groups
- ctx.save_for_backward(x, w)
-
gx = x.reshape(bs,ctx.groups,cin,x.shape[2],x.shape[3])
- str_x = np.lib.stride_tricks.as_strided(
- gx,
+ tx = np.lib.stride_tricks.as_strided(gx,
shape=(bs, ctx.groups, cin, oy, ox, H, W),
strides=(gx.strides[0], gx.strides[1], gx.strides[2], gx.strides[3]*ys, gx.strides[4]*xs, gx.strides[3], gx.strides[4]),
writeable=False,
)
- return np.einsum('igjYXyx,gkjyx -> igkYX', str_x, w.reshape(ctx.groups,rcout, cin, H, W)).reshape(bs, cout, oy, ox)
+ tw = w.reshape(ctx.groups, rcout, cin, H, W)
+ ctx.save_for_backward(tx, tw, x.shape)
+ return np.einsum('igjYXyx,gkjyx -> igkYX', tx, tw).reshape(bs, cout, oy, ox)
@staticmethod
def backward(ctx, grad_output):
bs,_,oy,ox = grad_output.shape
- x, w = ctx.saved_tensors
- cout,cin,H,W = w.shape
+ tx, tw, x_shape = ctx.saved_tensors
+ _,rcout,cin,H,W = tw.shape
ys,xs = ctx.stride
- rcout = cout//ctx.groups
-
- dx, dw = np.zeros_like(x), np.zeros_like(w)
-
- gdx = dx.reshape(bs,ctx.groups,cin,x.shape[2],x.shape[3])
- gdw = dw.reshape(ctx.groups,rcout, cin, H, W)
-
- ggg = grad_output.reshape(grad_output.shape[0],ctx.groups,rcout,oy,ox)
-
- gx = x.reshape(x.shape[0],ctx.groups,cin,x.shape[2],x.shape[3])
- str_x = np.lib.stride_tricks.as_strided(
- gx,
- shape=(bs, ctx.groups, cin, oy, ox, H, W),
- strides=(gx.strides[0], gx.strides[1], gx.strides[2], gx.strides[3]*ys, gx.strides[4]*xs, gx.strides[3], gx.strides[4]),
- writeable=False,
- )
+ OY,OX = x_shape[2:4]
- gdw = np.einsum('igkYX,igjYXyx -> gkjyx',ggg,str_x)
+ ggg = grad_output.reshape(bs,ctx.groups,rcout,oy,ox)
+ gdw = np.einsum('igkYX,igjYXyx -> gkjyx',ggg,tx)
#needs to be optimized
+ gdx = np.zeros((bs,ctx.groups,cin,OY,OX), dtype=tx.dtype)
for Y in range(grad_output.shape[2]):
for X in range(grad_output.shape[3]):
iY,iX = Y*ys, X*xs
- gdx[:,:,: , iY:iY+H, iX:iX+W] += np.einsum('igk,gkjyx->igjyx',ggg[:,:,:,Y,X], w.reshape(ctx.groups,rcout, cin, H, W))
+ gdx[:,:,: , iY:iY+H, iX:iX+W] += np.einsum('igk,gkjyx->igjyx',ggg[:,:,:,Y,X], tw)
- return gdx.reshape(x.shape), gdw.reshape(w.shape)
+ return gdx.reshape((bs, ctx.groups*cin, OY, OX)), gdw.reshape((ctx.groups*rcout, cin, H, W))
register('conv2d', Conv2D)

simpler conv2d