avgpool2d is a second class op
MODIFY
README.md
README.md
@@ -78,7 +78,7 @@ loss.backward()
optim.step()
```
-## GPU Support
+## GPU and Accelerator Support
tinygrad supports GPUs through PyOpenCL.
@@ -105,13 +105,15 @@ Warning: do not rely on the ANE port. It segfaults sometimes. So if you were doi
### Adding an accelerator
-You need to support 15 basic ops:
+You need to support 14 basic ops:
```
-Add, Sub, Mul, Pow, Sum, Dot
-Pad2D, Reshape, Transpose
-Relu, Log, Exp
-Conv2D, MaxPool2D, AvgPool2D
+Add, Sub, Mul, Pow # binary ops
+Dot # matrix multiplication
+Sum # reduce op
+Relu, Log, Exp # unary ops
+Conv2D, MaxPool2D # 2D ops
+Pad2D, Reshape, Transpose # moving things around ops
```
## ImageNet inference

MODIFY
test/test_ops.py
test/test_ops.py
@@ -169,7 +169,7 @@ class TestOps(unittest.TestCase):
with self.subTest(kernel_size=ksz):
helper_test_op([shape],
lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=ksz),
- lambda x: Tensor.avg_pool2d(x, kernel_size=ksz), device=self.device)
+ lambda x: Tensor.avg_pool2d(x, kernel_size=ksz), device=self.device, rtol=1e-5)
@unittest.skipUnless(GPU, "Requires GPU")
class TestOpsGPU(TestOps):

MODIFY
tinygrad/ops_cpu.py
tinygrad/ops_cpu.py
@@ -255,16 +255,3 @@ class MaxPool2D(Function):
return unstack_for_pool(lambda idx: grad_output * (idxs == idx), s, *ctx.kernel_size)
register('max_pool2d', MaxPool2D)
-class AvgPool2D(Function):
- @staticmethod
- def forward(ctx, x, kernel_size=(2, 2)):
- stack = stack_for_pool(x, *kernel_size)
- ctx.save_for_backward(x.shape)
- return np.mean(stack, axis=0)
-
- @staticmethod
- def backward(ctx, grad_output):
- s, = ctx.saved_tensors
- py, px = ctx.kernel_size
- return unstack_for_pool(lambda idx: grad_output/py/px, s, py, px)
-register('avg_pool2d', AvgPool2D)

MODIFY
tinygrad/ops_gpu.py
tinygrad/ops_gpu.py
@@ -389,21 +389,6 @@ class Exp(Function):
return binary_op(ctx, 'a * b', grad_output, ret)
register('exp', Exp, device=Device.GPU)
-class AvgPool2D(Function):
- @staticmethod
- def forward(ctx, input, kernel_size=(2, 2)):
- ret = subsample_op(ctx, input, kernel_size, kernel_size, iter_op="sumval += input[iid]",
- result_op="sumval / (ksz.x * ksz.y)", decls="float sumval=0.f")
- ctx.save_for_backward(input.shape)
- return ret
-
- @staticmethod
- def backward(ctx, grad_output):
- orig_shape, = ctx.saved_tensors
- return supersample_op(ctx, grad_output, orig_shape, ctx.kernel_size,
- result_op="input[iid] / (ksz.x * ksz.y)")
-register('avg_pool2d', AvgPool2D, device=Device.GPU)
-
class MaxPool2D(Function):
@staticmethod
def forward(ctx, input, kernel_size=(2, 2)):

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -241,6 +241,12 @@ class Tensor:
def abs(self):
return self.relu() + (-1.0*self).relu()
+ def avg_pool2d(self, kernel_size=(2,2)):
+ chan = self.shape[1]
+ ww = np.zeros((chan, 1, kernel_size[0], kernel_size[1]), dtype=np.float32)
+ ww[range(chan), 0, :, :] = 1/(kernel_size[0]*kernel_size[1])
+ return self.conv2d(Tensor(ww, device=self.device, requires_grad=False), stride=kernel_size, groups=chan)
+
# An instantiation of the Function is the Context
class Function:
def __init__(self, *tensors):

avgpool2d is a second class op