move batchnorm to Tensor
MODIFY
extra/onnx.py
extra/onnx.py
@@ -4,7 +4,6 @@ import functools
from onnx.mapping import TENSOR_TYPE_TO_NP_TYPE
from tinygrad.tensor import Tensor
from tinygrad.helpers import prod
-from tinygrad.nn import batch_normalize
from tinygrad.ops import DEBUG
def get_run_onnx(onnx_model):
@@ -114,7 +113,7 @@ def get_run_onnx(onnx_model):
ret = ret.reshape([s for i,s in enumerate(shape) if i != axis]) if len(indices) == 1 else ret # squeeze if needed
elif n.op_type == "BatchNormalization":
invstd = inp[4].add(opt.get('epsilon', 1e-5))**-0.5
- ret = batch_normalize(inp[0], inp[1], inp[2], inp[3], invstd)
+ ret = inp[0].batchnorm(inp[1], inp[2], inp[3], invstd)
elif n.op_type == "Gemm": ret = inp[0].linear(inp[1].transpose() if opt.get('transB', 0) == 1 else inp[1], inp[2])
elif n.op_type == "Conv":
x,w,b = inp if len(inp) == 3 else (inp[0], inp[1], None)

MODIFY
tinygrad/nn/__init__.py
tinygrad/nn/__init__.py
from tinygrad.tensor import Tensor
-def batch_normalize(x, weight, bias, mean, invstd):
- x = (x - mean.reshape(shape=[1, -1, 1, 1])) * weight.reshape(shape=[1, -1, 1, 1])
- return x.mul(invstd.reshape(shape=[1, -1, 1, 1])) + bias.reshape(shape=[1, -1, 1, 1])
-
class BatchNorm2D:
def __init__(self, sz, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1):
assert affine, "BatchNorm2D is only supported with affine"
@@ -38,7 +34,7 @@ class BatchNorm2D:
self.batch_invstd = batch_var.add(self.eps)**-0.5
batch_invstd = self.batch_invstd
- return batch_normalize(x, self.weight, self.bias, batch_mean, batch_invstd)
+ return x.batchnorm(self.weight, self.bias, batch_mean, batch_invstd)
# TODO: is this good weight init?
class Conv2d:

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -316,6 +316,10 @@ class Tensor:
y = (self - self.mean(axis=axis, keepdim=True))
return y.div((y*y).mean(axis=axis, keepdim=True).add(eps).sqrt())
+ def batchnorm(self, weight:Tensor, bias:Tensor, mean:Tensor, invstd:Tensor):
+ self = (self - mean.reshape(shape=[1, -1, 1, 1])) * weight.reshape(shape=[1, -1, 1, 1])
+ return self.mul(invstd.reshape(shape=[1, -1, 1, 1])) + bias.reshape(shape=[1, -1, 1, 1])
+
# An instantiation of the Function is the Context
class Function:
def __init__(self, device:str, *tensors:Tensor):

move batchnorm to Tensor