geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

bn: don't track_running_stats if you aren't training

Browse Source
This commit is contained in:
geohot
2022-06-05 16:41:20 -07:00
parent 1de75b67d5
commit 30025f2bc7
1 changed files with 6 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/nn.py
+6 -7
View File
@@ -12,19 +12,18 @@ class BatchNorm2D:
self.num_batches_tracked = Tensor.zeros(1, requires_grad=False)
def __call__(self, x):
if self.track_running_stats or Tensor.training:
if Tensor.training:
x_detached = x.detach()
batch_mean = x_detached.mean(axis=(0,2,3))
y = (x_detached - batch_mean.reshape(shape=[1, -1, 1, 1]))
batch_var = (y*y).mean(axis=(0,2,3))
if self.track_running_stats:
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
if self.num_batches_tracked is None: self.num_batches_tracked = Tensor.zeros(1, requires_grad=False)
self.num_batches_tracked += 1
if self.track_running_stats:
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
if self.num_batches_tracked is None: self.num_batches_tracked = Tensor.zeros(1, requires_grad=False)
self.num_batches_tracked += 1
if Tensor.training:
return self.normalize(x, batch_mean, batch_var)
return self.normalize(x, self.running_mean, self.running_var)