flip batchnorm function order
MODIFY
test/test_onnx.py
test/test_onnx.py
@@ -85,7 +85,7 @@ def run_onnx(onnx_model, inputs={}, debug=False):
ret = a.linear(w,b)
elif n.op_type == "BatchNormalization":
from tinygrad.nn import batch_normalize
- ret = batch_normalize(inp[0], inp[3], inp[4], inp[1], inp[2], opt['epsilon'])
+ ret = batch_normalize(inp[0], inp[1], inp[2], inp[3], inp[4], opt['epsilon'])
elif n.op_type == "MaxPool":
ret = inp[0].pad2d(opt['pads'])
ret = ret.max_pool2d(opt['kernel_shape'])
@@ -132,7 +132,7 @@ class TestOnnxModel(unittest.TestCase):
onnx_model = onnx.load(io.BytesIO(dat))
from test.test_efficientnet import chicken_img, car_img, _LABELS
inputs = {"data": chicken_img}
- tinygrad_out = run_onnx(onnx_model, inputs)['resnetv15_dense0_fwd'].numpy()
+ tinygrad_out = run_onnx(onnx_model, inputs, False)['resnetv15_dense0_fwd'].numpy()
cls = tinygrad_out.argmax()
print(cls, _LABELS[cls])

MODIFY
tinygrad/nn.py
tinygrad/nn.py
from tinygrad.tensor import Tensor
import numpy as np
-def batch_normalize(x, mean, var, weight, bias, eps):
+def batch_normalize(x, weight, bias, mean, var, eps):
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])
@@ -28,9 +28,9 @@ class BatchNorm2D:
if self.num_batches_tracked is None: self.num_batches_tracked = Tensor.zeros(1, requires_grad=False)
self.num_batches_tracked += 1
- return batch_normalize(x, batch_mean, batch_var, self.weight, self.bias, self.eps)
+ return batch_normalize(x, self.weight, self.bias, batch_mean, batch_var, self.eps)
- return batch_normalize(x, self.running_mean, self.running_var, self.weight, self.bias, self.eps)
+ return batch_normalize(x, self.weight, self.bias, self.running_mean, self.running_var,self.eps)
class Conv2d:
def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, bias=True):

flip batchnorm function order