way over the line limit
MODIFY
test/test_efficientnet.py
test/test_efficientnet.py
@@ -17,7 +17,7 @@ def _load_labels():
_LABELS = _load_labels()
-def _infer(model: EfficientNet, img):
+def _infer(model: EfficientNet, img, bs=1):
# preprocess image
aspect_ratio = img.size[0] / img.size[1]
img = img.resize((int(224*max(aspect_ratio,1.0)), int(224*max(1.0/aspect_ratio,1.0))))
@@ -34,8 +34,9 @@ def _infer(model: EfficientNet, img):
img /= np.array([0.229, 0.224, 0.225]).reshape((1, -1, 1, 1))
# run the net
+ if bs > 1: img = img.repeat(bs, axis=0)
out = model.forward(Tensor(img)).cpu()
- return _LABELS[np.argmax(out.data)]
+ return _LABELS[np.argmax(out.data[0])]
chicken_img = Image.open(pathlib.Path(__file__).parent / 'efficientnet/Chicken.jpg')
car_img = Image.open(pathlib.Path(__file__).parent / 'efficientnet/car.jpg')
@@ -54,6 +55,10 @@ class TestEfficientNet(unittest.TestCase):
label = _infer(self.model, chicken_img)
self.assertEqual(label, "hen")
+ def test_chicken_bigbatch(self):
+ label = _infer(self.model, chicken_img, 16)
+ self.assertEqual(label, "hen")
+
def test_car(self):
label = _infer(self.model, car_img)
self.assertEqual(label, "sports car, sport car")

MODIFY
tinygrad/helpers.py
tinygrad/helpers.py
@@ -20,3 +20,30 @@ def binary_broadcast(x_shape, y_shape, extra=False):
push(np.int32(max(shape_x[i], shape_y[i])), (shape_x[i] > 1, shape_y[i] > 1))
return (shape_ret, dimlist, complist) if extra else shape_ret
+
+from enum import Enum
+class UnaryOps(Enum):
+ RELU = 0
+ EXP = 1
+ LOG = 2
+ NEG = 3
+
+class BinaryOps(Enum):
+ ADD = 0
+ SUB = 1
+ MUL = 2
+ DIV = 3
+ POW = 4
+ A = 5
+ CMP = 6
+ # these are "complex" ops, and should probably be replaced
+ EXPMUL = 7
+ POW_D1 = 8
+ POW_D2 = 9
+ RELU_D = 10
+
+class ReduceOps(Enum):
+ SUM = 0
+ MAX = 1
+ # this is "complex"
+ NEGSUM = 2
\ No newline at end of file

MODIFY
tinygrad/llops/opencl.py
tinygrad/llops/opencl.py
import functools
import numpy as np
import pyopencl as cl
-from tinygrad.helpers import binary_broadcast
+from tinygrad.helpers import binary_broadcast, UnaryOps, BinaryOps, ReduceOps
cl_ctx, cl_queue = None, None
def require_init_gpu():
@@ -48,7 +48,12 @@ def clbuild(name, prg):
def run(*args): clprg(cl_queue, *args)
return run
-def unary_op(code, x, ret):
+def unary_op(op, x, ret):
+ if op == UnaryOps.RELU: code = 'max(a, (float)0.)'
+ elif op == UnaryOps.EXP: code = 'exp(a)'
+ elif op == UnaryOps.LOG: code = 'log(a)'
+ elif op == UnaryOps.NEG: code = '-a'
+ else: raise Exception(f"{op} isn't supported")
unop = clbuild("unop", """
__kernel void unop(__global const float4 *a_g, __global float4 *res_g) {
int gid = get_global_id(0);
@@ -78,16 +83,33 @@ def get_binop_prg(code, complist):
res_g[gid0] = """+code+""";\n}"""
return cl.Program(cl_ctx, prg).build(), dtype[2] == "float4"
-def binary_op(code, x, y, ret):
+def binary_op(op, x, y, ret):
+ if op == BinaryOps.ADD: code = "a+b"
+ elif op == BinaryOps.SUB: code = "a-b"
+ elif op == BinaryOps.MUL: code = "a*b"
+ elif op == BinaryOps.DIV: code = "b/a"
+ elif op == BinaryOps.POW: code = "pow(a,b)"
+ elif op == BinaryOps.A: code = "a"
+ elif op == BinaryOps.CMP: code = "1.0f*(a==b)"
+ elif op == BinaryOps.EXPMUL: code = "exp(a) * b"
+ elif op == BinaryOps.POW_D1: code = 'b * pow(a, b-1.0f)'
+ elif op == BinaryOps.POW_D2: code = 'log(a) * pow(a, b)'
+ elif op == BinaryOps.RELU_D: code = 'b * max(sign(a), (float)0.)'
+ else: raise Exception(f"{op} isn't supported")
+
shape_ret, dimlist, complist = binary_broadcast(x.shape, y.shape, True)
assert tuple(shape_ret) == tuple(ret.shape)
prod_list = np.array(dimlist, dtype=i32)[-1::-1].cumprod(dtype=i32)[-1::-1] # take cumprod from back to front
prg, is_float4 = get_binop_prg(code, tuple(complist))
- kernel_size = [(roundup(prod_list[0])//4) if is_float4 else prod_list[0]] if len(dimlist) > 0 else [1]
- prg.binop(cl_queue, kernel_size, None, x.cl, y.cl, ret.cl, *dimlist, *(prod_list[1:]))
+ kernel_size = ((roundup(prod_list[0])//4) if is_float4 else prod_list[0]) if len(dimlist) > 0 else 1
+ prg.binop(cl_queue, [kernel_size], None, x.cl, y.cl, ret.cl, *dimlist, *(prod_list[1:]))
return ret
-def reduce_op(code, inp, ret, start="0.0"):
+def reduce_op(op, inp, ret, start="0.0"):
+ if op == ReduceOps.SUM: code = "out += a"
+ elif op == ReduceOps.MAX: code = "out = max(a,out)"
+ elif op == ReduceOps.NEGSUM: code = "out -= a"
+ else: raise Exception(f"{op} isn't supported")
# TODO: this is insanely slow
# NOTE: ret.shape can be (1,), it's mostly by luck that this works
reduce = clbuild("reduce", """

MODIFY
tinygrad/ops/ops_gpu.py
tinygrad/ops/ops_gpu.py
import numpy as np
-from tinygrad.helpers import binary_broadcast
+from tinygrad.helpers import binary_broadcast, UnaryOps, BinaryOps, ReduceOps
from ..tensor import Function
from ..llops.opencl import GPUBuffer as Buffer
from ..llops.opencl import unary_op, binary_op, reduce_op, perm_axis, inner_slice
@@ -17,16 +17,16 @@ class UnaryOp(Function):
return binary_op(ctx.bop, input, grad_output, Buffer(input.shape))
class ReLU(UnaryOp):
- fop = 'max(a, (float)0.)'
- bop = 'b * max(sign(a), (float)0.)'
+ fop = UnaryOps.RELU
+ bop = BinaryOps.RELU_D
class Log(UnaryOp):
- fop = 'log(a)'
- bop = 'b / a'
+ fop = UnaryOps.LOG
+ bop = BinaryOps.DIV
class Exp(UnaryOp):
- fop = 'exp(a)'
- bop = 'b * exp(a)'
+ fop = UnaryOps.EXP
+ bop = BinaryOps.EXPMUL
# ************* reduce ops *************
@@ -36,36 +36,36 @@ def reduce_shape(shape, axis):
class Sum(Function):
def forward(ctx, input, axis=None):
ctx.save_for_backward(input.shape)
- return reduce_op("out += a", input, Buffer(reduce_shape(input.shape, axis)))
+ return reduce_op(ReduceOps.SUM, input, Buffer(reduce_shape(input.shape, axis)))
def backward(ctx, grad_output):
shape_input, = ctx.saved_tensors
# NOTE: the b Buffer isn't used, since this is just for broadcast
ret = Buffer(shape_input)
- return binary_op('a', grad_output, ret, ret)
+ return binary_op(BinaryOps.A, grad_output, ret, ret)
class Max(Function):
def forward(ctx, input, axis=None):
- ret = reduce_op("out = max(a,out)", input, Buffer(reduce_shape(input.shape, axis)), start="-INFINITY")
+ ret = reduce_op(ReduceOps.MAX, input, Buffer(reduce_shape(input.shape, axis)), start="-INFINITY")
ctx.save_for_backward(input, ret)
return ret
def backward(ctx, grad_output):
input, ret = ctx.saved_tensors
- ret2 = binary_op("1.0f*(a==b)", input, ret, Buffer(input.shape))
- div = reduce_op("out += a", ret2, Buffer(grad_output.shape), start="1e-10")
- binary_op("a/b", ret2, div, ret2)
- return binary_op('a*b', ret2, grad_output, ret2)
+ ret2 = binary_op(BinaryOps.CMP, input, ret, Buffer(input.shape))
+ div = reduce_op(ReduceOps.SUM, ret2, Buffer(grad_output.shape), start="1e-10")
+ binary_op(BinaryOps.DIV, div, ret2, ret2)
+ return binary_op(BinaryOps.MUL, ret2, grad_output, ret2)
# ************* binary ops *************
-def unbroadcast(out, in_sh, code="out += a"):
- return reduce_op(code, out, Buffer(in_sh))
+def unbroadcast(out, in_sh, op=ReduceOps.SUM):
+ return reduce_op(op, out, Buffer(in_sh))
class Add(Function):
def forward(ctx, x, y):
ctx.save_for_backward(x.shape, y.shape)
- return binary_op('a+b', x, y, Buffer(binary_broadcast(x.shape, y.shape)))
+ return binary_op(BinaryOps.ADD, x, y, Buffer(binary_broadcast(x.shape, y.shape)))
def backward(ctx, grad_output):
shape_x, shape_y = ctx.saved_tensors
@@ -75,35 +75,35 @@ class Add(Function):
class Sub(Function):
def forward(ctx, x, y):
ctx.save_for_backward(x.shape, y.shape)
- return binary_op('a-b', x, y, Buffer(binary_broadcast(x.shape, y.shape)))
+ return binary_op(BinaryOps.SUB, x, y, Buffer(binary_broadcast(x.shape, y.shape)))
def backward(ctx, grad_output):
shape_x, shape_y = ctx.saved_tensors
return unbroadcast(grad_output, shape_x) if ctx.needs_input_grad[0] else None, \
- unbroadcast(grad_output, shape_y, code="out -= a") if ctx.needs_input_grad[1] else None
+ unbroadcast(grad_output, shape_y, ReduceOps.NEGSUM) if ctx.needs_input_grad[1] else None
class Mul(Function):
def forward(ctx, x, y):
ctx.save_for_backward(x, y)
- return binary_op('a*b', x, y, Buffer(binary_broadcast(x.shape, y.shape)))
+ return binary_op(BinaryOps.MUL, x, y, Buffer(binary_broadcast(x.shape, y.shape)))
def backward(ctx, grad_output):
x,y = ctx.saved_tensors
tmp = Buffer(grad_output.shape)
- grad_x = unbroadcast(binary_op('a*b', y, grad_output, tmp), x.shape) if ctx.needs_input_grad[0] else None
- grad_y = unbroadcast(binary_op('a*b', x, grad_output, tmp), y.shape) if ctx.needs_input_grad[1] else None
+ grad_x = unbroadcast(binary_op(BinaryOps.MUL, y, grad_output, tmp), x.shape) if ctx.needs_input_grad[0] else None
+ grad_y = unbroadcast(binary_op(BinaryOps.MUL, x, grad_output, tmp), y.shape) if ctx.needs_input_grad[1] else None
return grad_x, grad_y
class Pow(Function):
def forward(ctx, x, y):
ctx.save_for_backward(x, y)
- return binary_op('pow(a,b)', x, y, Buffer(binary_broadcast(x.shape, y.shape)))
+ return binary_op(BinaryOps.POW, x, y, Buffer(binary_broadcast(x.shape, y.shape)))
def backward(ctx, grad_output):
x,y = ctx.saved_tensors
tmp = Buffer(grad_output.shape)
- grad_x = unbroadcast(binary_op('a*b', grad_output, binary_op('b * pow(a, b-1.0f)', x, y, tmp), tmp), x.shape) if ctx.needs_input_grad[0] else None
- grad_y = unbroadcast(binary_op('a*b', grad_output, binary_op('log(a) * pow(a, b)', x, y, tmp), tmp), y.shape) if ctx.needs_input_grad[1] else None
+ grad_x = unbroadcast(binary_op(BinaryOps.MUL, grad_output, binary_op(BinaryOps.POW_D1, x, y, tmp), tmp), x.shape) if ctx.needs_input_grad[0] else None
+ grad_y = unbroadcast(binary_op(BinaryOps.MUL, grad_output, binary_op(BinaryOps.POW_D2, x, y, tmp), tmp), y.shape) if ctx.needs_input_grad[1] else None
return grad_x, grad_y
# ************* movement ops *************

way over the line limit