From 6842ad9ec8a0b8ed512e26cfc00c6dcbb20dc9b7 Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Sun, 3 Jan 2021 08:14:16 -0800
Subject: [PATCH] minor cleanups, yolo work

---
 .github/workflows/test.yml |  2 +-
 examples/yolov5.py         | 32 +++++++++++++++++---------------
 tinygrad/optim.py          |  8 +++-----
 3 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 04577023..0ffa74ab 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -13,7 +13,7 @@ jobs:
     - name: Install SLOCCount
       run: sudo apt-get install sloccount
     - name: Check <1000 lines
-      run: sloccount tinygrad test examples; if [ $(sloccount tinygrad | sed -n 's/.*Total Physical Source Lines of Code (SLOC)[ ]*= \([^ ]*\).*/\1/p' | tr -d ',') -gt 1000 ]; then exit 1; fi
+      run: sloccount tinygrad test examples extra; if [ $(sloccount tinygrad | sed -n 's/.*Total Physical Source Lines of Code (SLOC)[ ]*= \([^ ]*\).*/\1/p' | tr -d ',') -gt 1000 ]; then exit 1; fi
 
   linter:
     name: Indentation  Linter
diff --git a/examples/yolov5.py b/examples/yolov5.py
index db48c762..6f95ed5f 100755
--- a/examples/yolov5.py
+++ b/examples/yolov5.py
@@ -5,27 +5,29 @@ from extra.utils import fetch, my_unpickle
 
 if __name__ == "__main__":
   dat = fetch('https://github.com/ultralytics/yolov5/releases/download/v3.0/yolov5s.pt')
-  #import torch
-  #td = torch.load(io.BytesIO(dat))
-  #print(td)
+  #dat = fetch('https://github.com/ultralytics/yolov5/releases/download/v3.0/yolov5m.pt')
 
   import zipfile
   fp = zipfile.ZipFile(io.BytesIO(dat))
   #fp.printdir()
   data = fp.read('archive/data.pkl')
 
-  #import pickletools
-  #pickletools.dis(io.BytesIO(data))
-
+  # yolo specific
   ret, out = my_unpickle(io.BytesIO(data))
-  print(dir(ret['model']))
-  for m in ret['model']._modules['model']:
-    print(m)
-    print(m._modules.keys())
+  d = ret['model'].yaml
+  for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']):
+    tm = ret['model']._modules['model'][i]
+    print(i, f, n, m, args, tm._modules.keys())
+    # Focus, Conv, BottleneckCSP, SPP, Concat, Detect
+    #for k,v in tm._modules.items():
+    #  print("   ", k, v)
+    if m in "Focus":
+      conv = tm._modules['conv']
+      print("   ", conv._modules)
+    if m in "Conv":
+      conv, bn = tm._modules['conv'], tm._modules['bn']
+      print("   ", conv)
+      #print(bn)
+
 
-  """
-  weights = fake_torch_load(data)
-  for k,v in weights:
-    print(k)
-  """
 
diff --git a/tinygrad/optim.py b/tinygrad/optim.py
index 722298b2..49f052aa 100644
--- a/tinygrad/optim.py
+++ b/tinygrad/optim.py
@@ -1,6 +1,4 @@
 # sorted in order of increasing complexity
-
-import numpy as np
 from tinygrad.tensor import Tensor
 
 class Optimizer:
@@ -25,7 +23,7 @@ class RMSprop(Optimizer):
     super().__init__(params)
     self.lr, self.decay, self.eps = lr, decay, eps
 
-    self.v = [Tensor(np.zeros(t.shape, dtype=np.float32), device=params[0].device, requires_grad=False) for t in self.params]
+    self.v = [Tensor.zeros(*t.shape, device=params[0].device, requires_grad=False) for t in self.params]
 
   def step(self):
     for i, t in enumerate(self.params):
@@ -37,8 +35,8 @@ class Adam(Optimizer):
     super().__init__(params)
     self.lr, self.b1, self.b2, self.eps, self.t = lr, b1, b2, eps, 0
 
-    self.m = [Tensor(np.zeros(t.shape, dtype=np.float32), device=params[0].device, requires_grad=False) for t in self.params]
-    self.v = [Tensor(np.zeros(t.shape, dtype=np.float32), device=params[0].device, requires_grad=False) for t in self.params]
+    self.m = [Tensor.zeros(*t.shape, device=params[0].device, requires_grad=False) for t in self.params]
+    self.v = [Tensor.zeros(*t.shape, device=params[0].device, requires_grad=False) for t in self.params]
 
   def step(self):
     self.t = self.t + 1