fix bn training
MODIFY
examples/train_efficientnet.py
examples/train_efficientnet.py
@@ -5,6 +5,7 @@ from models.efficientnet import EfficientNet
from tinygrad.tensor import Tensor
from extra.utils import get_parameters
from tqdm import trange
+from tinygrad.nn import BatchNorm2D
import tinygrad.optim as optim
from datasets import fetch_cifar
@@ -13,12 +14,16 @@ class TinyConvNet:
conv = 3
inter_chan, out_chan = 8, 16 # for speed
self.c1 = Tensor.uniform(inter_chan,3,conv,conv)
+ self.bn1 = BatchNorm2D(inter_chan)
self.c2 = Tensor.uniform(out_chan,inter_chan,conv,conv)
+ self.bn2 = BatchNorm2D(out_chan)
self.l1 = Tensor.uniform(out_chan*6*6, classes)
def forward(self, x):
x = x.conv2d(self.c1).relu().max_pool2d()
+ x = self.bn1(x)
x = x.conv2d(self.c2).relu().max_pool2d()
+ x = self.bn2(x)
x = x.reshape(shape=[x.shape[0], -1])
return x.dot(self.l1).logsoftmax()
@@ -26,7 +31,6 @@ if __name__ == "__main__":
X_train, Y_train = fetch_cifar()
classes = 10
- Tensor.default_gpu = os.getenv("GPU") is not None
TINY = os.getenv("TINY") is not None
TRANSFER = os.getenv("TRANSFER") is not None
if TINY:
@@ -38,11 +42,11 @@ if __name__ == "__main__":
model = EfficientNet(int(os.getenv("NUM", "0")), classes, has_se=False)
parameters = get_parameters(model)
- print("parameters", len(parameters))
+ print("parameter count", len(parameters))
optimizer = optim.Adam(parameters, lr=0.001)
- #BS, steps = 16, 32
- BS, steps = 64 if TINY else 16, 2048
+ BS, steps = int(os.getenv("BS", "64" if TINY else "16")), 2048
+ print("training with batch size %d for %d steps" % (BS, steps))
for i in (t := trange(steps)):
samp = np.random.randint(0, X_train.shape[0], size=(BS))

MODIFY
tinygrad/nn.py
tinygrad/nn.py
@@ -2,8 +2,9 @@ from tinygrad.tensor import Tensor
import numpy as np
class BatchNorm2D:
- def __init__(self, sz, eps=1e-5, track_running_stats=False, training=False, momentum=0.1):
- self.eps, self.track_running_stats, self.training, self.momentum = eps, track_running_stats, training, momentum
+ def __init__(self, sz, eps=1e-5, track_running_stats=False, momentum=0.1):
+ # TODO: add affine?
+ self.eps, self.track_running_stats, self.momentum = eps, track_running_stats, momentum
self.weight, self.bias = Tensor.ones(sz), Tensor.zeros(sz)
@@ -11,7 +12,7 @@ class BatchNorm2D:
self.num_batches_tracked = Tensor.zeros(1, requires_grad=False)
def __call__(self, x):
- if self.track_running_stats or self.training:
+ if self.track_running_stats or Tensor.training:
batch_mean = x.mean(axis=(0,2,3))
y = (x - batch_mean.reshape(shape=[1, -1, 1, 1]))
batch_var = (y*y).mean(axis=(0,2,3))
@@ -22,7 +23,7 @@ class BatchNorm2D:
if self.num_batches_tracked is None: self.num_batches_tracked = Tensor.zeros(1, requires_grad=False)
self.num_batches_tracked += 1
- if self.training:
+ if Tensor.training:
return self.normalize(x, batch_mean, batch_var)
return self.normalize(x, self.running_mean, self.running_var)

fix bn training