github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-23 05:48:08 -05:00
Code Issues Packages Projects Releases Wiki Activity

track running stats by default, and detach

Browse Source
This commit is contained in:
George Hotz
2022-01-15 21:46:54 -08:00
parent d541e2a8e5
commit d844ecee27
2 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
tinygrad/nn.py
View File

@@ -2,7 +2,7 @@ from tinygrad.tensor import Tensor
import numpy as np
class BatchNorm2D:
def __init__(self, sz, eps=1e-5, affine=True, track_running_stats=False, momentum=0.1):
def __init__(self, sz, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1):
assert affine == True
self.eps, self.track_running_stats, self.momentum = eps, track_running_stats, momentum
@@ -13,8 +13,9 @@ class BatchNorm2D:
def __call__(self, x):
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]))
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:

2
tinygrad/tensor.py
View File

@@ -169,7 +169,7 @@ class Tensor:
return ret
def detach(self):
return Tensor(self.data, device=self.device)
return Tensor(self.data, device=self.device, requires_grad=False)
# ***** non first class ops *****