From f9788eba144114bacb09798e669b38279400fa24 Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Tue, 27 Oct 2020 08:53:35 -0700
Subject: [PATCH] parameters, and start on efficientnet

---
 test/test_mnist.py              | 12 +++++--
 tinygrad/models/efficientnet.py | 57 +++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 3 deletions(-)
 create mode 100644 tinygrad/models/efficientnet.py

diff --git a/test/test_mnist.py b/test/test_mnist.py
index 6e32956f7d..f73a58f56b 100644
--- a/test/test_mnist.py
+++ b/test/test_mnist.py
@@ -16,6 +16,9 @@ class TinyBobNet:
     self.l1 = Tensor(layer_init_uniform(784, 128))
     self.l2 = Tensor(layer_init_uniform(128, 10))
 
+  def parameters(self):
+    return [self.l1, self.l2]
+
   def forward(self, x):
     return x.dot(self.l1).relu().dot(self.l2).logsoftmax()
 
@@ -30,6 +33,9 @@ class TinyConvNet:
     self.c2 = Tensor(layer_init_uniform(out_chan,inter_chan,conv,conv))
     self.l1 = Tensor(layer_init_uniform(out_chan*5*5, 10))
 
+  def parameters(self):
+    return [self.l1, self.c1, self.c2]
+
   def forward(self, x):
     x.data = x.data.reshape((-1, 1, 28, 28)) # hacks
     x = x.conv2d(self.c1).relu().max_pool2d()
@@ -80,21 +86,21 @@ class TestMNIST(unittest.TestCase):
   def test_conv(self):
     np.random.seed(1337)
     model = TinyConvNet()
-    optimizer = optim.Adam([model.c1, model.c2, model.l1], lr=0.001)
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
     train(model, optimizer, steps=200)
     evaluate(model)
     
   def test_sgd(self):
     np.random.seed(1337)
     model = TinyBobNet()
-    optimizer = optim.SGD([model.l1, model.l2], lr=0.001)
+    optimizer = optim.SGD(model.parameters(), lr=0.001)
     train(model, optimizer, steps=1000)
     evaluate(model)
     
   def test_rmsprop(self):
     np.random.seed(1337)
     model = TinyBobNet()
-    optimizer = optim.RMSprop([model.l1, model.l2], lr=0.0002)
+    optimizer = optim.RMSprop(model.parameters(), lr=0.0002)
     train(model, optimizer, steps=1000)
     evaluate(model)
 
diff --git a/tinygrad/models/efficientnet.py b/tinygrad/models/efficientnet.py
new file mode 100644
index 0000000000..a68089a488
--- /dev/null
+++ b/tinygrad/models/efficientnet.py
@@ -0,0 +1,57 @@
+# TODO: implement BatchNorm2d and Swish
+# aka batch_norm, pad, swish, dropout
+# https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b0-355c32eb.pth
+# a rough copy of
+# https://github.com/lukemelas/EfficientNet-PyTorch/blob/master/efficientnet_pytorch/model.py
+
+class BatchNorm2D:
+  def __init__(self, sz):
+    self.weight = Tensor.zeros(sz)
+    self.bias = Tensor.zeros(sz)
+    # TODO: need running_mean and running_var
+
+  def __call__(self, x):
+    # this work at inference?
+    return x * self.weight + self.bias
+
+class MBConvBlock:
+  def __init__(self, d0, d1, d2, d3):
+    self._expand_conv = Tensor.zeros(d1, d0, 1, 1)
+    self._bn0 = BatchNorm2D(d1)
+    self._depthwise_conv = Tensor.zeros(d1, 1, 3, 3)
+    self._bn1 = BatchNorm2D(d1)
+    self._se_reduce = Tensor.zeros(d2, d1, 1, 1)
+    self._se_reduce_bias = Tensor.zeros(d2)
+    self._se_expand = Tensor.zeros(d1, d2, 1, 1)
+    self._se_expand_bias = Tensor.zeros(d1)
+    self._project_conv = Tensor.zeros(d3, d2, 1, 1)
+    self._bn2 = BatchNorm2D(d3)
+
+  def __call__(self, x):
+    x = self._bn0(x.conv2d(self._expand_conv))
+    x = self._bn1(x.conv2d(self._depthwise_conv))  # TODO: repeat on axis 1
+    x = x.conv2d(self._se_reduce) + self._se_reduce_bias
+    x = x.conv2d(self._se_expand) + self._se_expand_bias
+    x = self._bn2(x.conv2d(self._project_conv))
+    return x.swish()
+
+class EfficientNet:
+  def __init__(self):
+    self._conv_stem = Tensor.zeros(32, 3, 3, 3)
+    self._bn0 = BatchNorm2D(32)
+    self._blocks = []
+    # TODO: create blocks
+
+    self._conv_head = Tensor.zeros(1280, 320, 1, 1)
+    self._bn1 = BatchNorm2D(1280)
+    self._fc = Tensor.zeros(1280, 1000)
+
+  def forward(x):
+    x = self._bn0(x.pad(0,1,0,1).conv2d(self._conv_stem, stride=2))
+    for b in self._blocks:
+      x = b(x)
+    x = self._bn1(x.conv2d(self._conv_head))
+    x = x.avg_pool2d() # wrong
+    x = x.dropout(0.2)
+    return x.dot(self_fc).swish()
+