gpu SGD runs now
MODIFY
tinygrad/opsgpu.py
tinygrad/opsgpu.py
@@ -31,10 +31,43 @@ class Add(Function):
return grad_output, grad_output
register('add', Add, gpu=True)
+class Sub(Function):
+ @staticmethod
+ def forward(ctx, x, y):
+ ret = buffer_like(ctx, x)
+ prg = cl.Program(ctx.cl_ctx, """
+ __kernel void sub(
+ __global const float *a_g, __global const float *b_g, __global float *res_g)
+ {
+ int gid = get_global_id(0);
+ res_g[gid] = a_g[gid] - b_g[gid];
+ }
+ """).build()
+ prg.sub(ctx.cl_queue, [np.prod(ret.shape)], None, x, y, ret)
+ return ret
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ return grad_output, grad_output
+register('sub', Sub, gpu=True)
+
class Mul(Function):
@staticmethod
def forward(ctx, x, y):
ret = buffer_like(ctx, x)
+
+ # HACK
+ if y.shape == (1,):
+ prg = cl.Program(ctx.cl_ctx, """
+ __kernel void mul(
+ __global const float *a_g, __global const float *b_g, __global float *res_g)
+ {
+ int gid = get_global_id(0);
+ res_g[gid] = a_g[gid] * b_g[0];
+ }
+ """).build()
+ prg.mul(ctx.cl_queue, [np.prod(ret.shape)], None, x, y, ret)
+ elif x.shape == y.shape:
prg = cl.Program(ctx.cl_ctx, """
__kernel void mul(
__global const float *a_g, __global const float *b_g, __global float *res_g)
@@ -44,6 +77,9 @@ class Mul(Function):
}
""").build()
prg.mul(ctx.cl_queue, [np.prod(ret.shape)], None, x, y, ret)
+ else:
+ raise Exception("mismatched shapes %r %r" % (x.shape, y.shape))
+
ctx.save_for_backward(x, y, prg)
return ret

MODIFY
tinygrad/optim.py
tinygrad/optim.py
@@ -10,7 +10,7 @@ class Optimizer:
class SGD(Optimizer):
def __init__(self, params, lr=0.001):
super(SGD, self).__init__(params)
- self.lr = Tensor([lr])
+ self.lr = Tensor([lr], gpu=params[0].gpu)
def step(self):
for t in self.params:

gpu SGD runs now