adding batch norm

Author: yunlongdong

cnnumpy/io.py

@@ -1,5 +1,6 @@
 import json, re, numpy as np
 from .net import Net
+import os
 
 def read_net(path):
 	net = Net()
@@ -38,7 +39,7 @@ def parse(matched):
 res = (flatten, upsample, conv, relu, sigmoid, maxpool, dense, concat, add, batchnorm, weight)
 
 def read_onnx(path):
-	with open(path+'.onnx') as f:
+	with open(path+'.txt') as f:
 		cont = f.read()
 	for i in res: cont = i.sub(parse, cont)
 	#for i in cont.split('\n'): print(i)
@@ -97,5 +98,9 @@ def read_onnx(path):
 
 	net = Net()
 	net.load_json(body, flow)
-	net.load_weights(np.load(path+'.npy'))
+	if os.path.exists(path+'_bn.npy'):
+		bn = np.load(path+'_bn.npy')
+	else:
+		bn = []
+	net.load_weights(np.load(path+'.npy'), bn)
 	return net

cnnumpy/layer.py

@@ -21,7 +21,7 @@ class Dense(Layer):
     name = 'dense'
     def __init__(self, c, n):
         self.K = np.zeros((n, c), dtype=np.float32)
-        self.bias = np.zeros(c, dtype=np.float32)
+        self.bias = np.zeros(n, dtype=np.float32)
 
     def para(self): return self.K.shape
 
@@ -91,6 +91,7 @@ def forward(self, x):
 class Maxpool(Layer):
     name = 'maxpool'
     def __init__(self, w=2, stride=2):
+        self.w = w
         self.stride = stride
 
     def para(self): return (self.stride,)
@@ -114,16 +115,32 @@ def __init__(self): pass
 
     def forward(self, x):
         return np.concatenate(x, axis=1)
-
+        
 class BatchNorm(Layer):
     name = 'batchnorm'
     def __init__(self, c):
         self.c = c
+        self.k = np.zeros(c, dtype=np.float32)
+        self.b = np.zeros(c, dtype=np.float32)
+        self.m = np.zeros(c, dtype=np.float32)
+        self.v = np.zeros(c, dtype=np.float32)
 
     def forward(self, x):
+
+        x = (x - self.m.reshape(1, -1, 1, 1))/np.sqrt(self.v.reshape(1, -1, 1, 1)) * self.k.reshape(1, -1, 1, 1) + self.b.reshape(1, -1, 1, 1)
+
         return x
 
     def load(self, buf):
+        c = self.c
+        self.k = buf[:c]
+        self.b = buf[c:2*c]
+        return self.c * 2
+
+    def load_bn(self, buf):
+        c = self.c
+        self.m = buf[:c]
+        self.v = buf[c:2*c]
         return self.c * 2
 
 layerkey = {'dense':Dense, 'conv':Conv2d, 'relu':ReLU, 'batchnorm':BatchNorm,

cnnumpy/net.py

@@ -66,10 +66,16 @@ def flw2code(self, style='list'):
 				body.append('')
 		return '\n'.join(body)
 
-	def load_weights(self, data):
+	def load_weights(self, data, data_bn):
 		s = 0
+		b = 0
 		for i in self.body: 
+			if 'batchnorm' in i[0]:
+				# print(i[0])
+				b += i[1].load_bn(data_bn[b:])
 			s += i[1].load(data[s:])
+		# print(data.shape, s)
+		# print(data_bn.shape, b)
 
 	def __call__(self, x):
 		return self.forward(x)