fix batchnorm folding in openpilot compile
MODIFY
extra/onnx.py
extra/onnx.py
@@ -81,7 +81,9 @@ def get_run_onnx(onnx_model):
elif n.op_type == "Transpose": ret = inp[0].permute(order=opt['perm'])
elif n.op_type == "Squeeze": ret = inp[0].reshape([s for i,s in enumerate(inp[0].shape) if i not in opt['axes']])
elif n.op_type == "GlobalAveragePool": ret = inp[0].mean(axis=tuple(range(2, len(inp[0].shape))), keepdim=True)
- elif n.op_type == "BatchNormalization": ret = batch_normalize(inp[0], inp[1], inp[2], inp[3], inp[4], opt.get('epsilon', 1e-5))
+ 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)
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
openpilot/compile.py
openpilot/compile.py
@@ -80,6 +80,14 @@ def compile(input, output_fn):
et = time.monotonic()
print(f"ran openpilot model in {(et-st)*1000.0:.2f} ms, waited {(mt2-mt)*1000.0:.2f} ms for realize, {(et-mt2)*1000.0:.2f} ms for GPU queue")
+ # realize all non GCed tensors (fix for batchnorm folding)
+ import gc
+ gc.collect()
+ def tensors_allocated():
+ return [x for x in gc.get_objects() if isinstance(x, Tensor)]
+ for x in tensors_allocated():
+ x.realize()
+
# real run
inputs, np_inputs = get_random_input_tensors()
tinygrad_out = run_onnx(inputs)['outputs']

MODIFY
tinygrad/nn/__init__.py
tinygrad/nn/__init__.py
from tinygrad.tensor import Tensor
-def batch_normalize(x, weight, bias, mean, var, eps):
+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(var.add(eps).reshape(shape=[1, -1, 1, 1])**-0.5) + bias.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):
@@ -18,7 +18,7 @@ class BatchNorm2D:
if Tensor.training:
# This requires two full memory accesses to x
# https://github.com/pytorch/pytorch/blob/c618dc13d2aa23625cb0d7ada694137532a4fa33/aten/src/ATen/native/cuda/Normalization.cuh
- # There's "online" algorithms that fix this
+ # There's "online" algorithms that fix this, like https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_Online_algorithm
x_detached = x.detach()
batch_mean = x_detached.mean(axis=(0,2,3))
y = (x_detached - batch_mean.reshape(shape=[1, -1, 1, 1]))
@@ -29,10 +29,12 @@ class BatchNorm2D:
self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * batch_mean
self.running_var = (1 - self.momentum) * self.running_var + self.momentum * batch_var
self.num_batches_tracked += 1
+ else:
+ batch_mean, batch_var = self.running_mean, self.running_var
- return batch_normalize(x, self.weight, self.bias, batch_mean, batch_var, self.eps)
-
- return batch_normalize(x, self.weight, self.bias, self.running_mean, self.running_var, self.eps)
+ # NOTE: this can be precomputed for static inference. if you manually update running_var, you have to reset this
+ if Tensor.training or getattr(self, "batch_invstd", None) is None: self.batch_invstd = batch_var.add(self.eps)**-0.5
+ return batch_normalize(x, self.weight, self.bias, batch_mean, self.batch_invstd)
class Conv2d:
def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, bias=True):

fix batchnorm folding in openpilot compile