load weights

Author: yxdragon

cnnumpy/layer.py

@@ -12,6 +12,8 @@ def backward(self, grad_y): pass
 
     def para(self): return None
 
+    def load(self, buf): return 0
+
     def __call__(self, x):
         return self.forward(x)
 
@@ -28,9 +30,11 @@ def forward(self, x):
         y += self.bias.reshape((1, -1))
         return y
 
-    def load(self, weights, bias):
-        self.K = weights.copy()
-        self.bias = bias.copy()
+    def load(self, buf):
+        sk, sb = self.K.size, self.bias.size
+        self.K.ravel()[:] = buf[:sk]
+        self.bias.ravel()[:] = buf[sk:sk+sb]
+        return sk + sb
 
 class Conv2d(Layer):
     name = 'conv'
@@ -46,9 +50,11 @@ def forward(self, x):
         out += self.bias.reshape((1, -1, 1, 1))
         return out
 
-    def load(self, weights, bias):
-        self.K = weights.copy()
-        self.bias = bias.copy()
+    def load(self, buf):
+        sk, sb = self.K.size, self.bias.size
+        self.K.ravel()[:] = buf[:sk]
+        self.bias.ravel()[:] = buf[sk:sk+sb]
+        return sk + sb
 
 class ReLU(Layer):
     name = 'relu'

cnnumpy/net.py

@@ -67,18 +67,10 @@ def flw2code(self, style='list'):
 				body.append('')
 		return '\n'.join(body)
 
-	def load_data(self, data, shpkey):
+	def load_weights(self, data):
 		s = 0
-		for i in shpkey:
-			for j in range(len(shpkey[i])):
-				sp = shpkey[i][j]
-				l = np.cumprod(sp)[-1]
-				shpkey[i][j] = data[s:s+l].reshape(sp)
-				s += l
-		dic = dict(self.body)
-		for i in shpkey:
-			dic[i].K = shpkey[i][0]
-			dic[i].bias = shpkey[i][1]
+		for i in self.body: 
+			s += i[1].load(data[s:])
 
 	def __call__(self, x):
 		return self.forward(x)
@@ -88,10 +80,8 @@ def read_net(path):
 	with open(path+'.lay') as f: lay = json.load(f)
 	with open(path+'.flw') as f: flw = json.load(f)
 	net.load_json(lay, flw)
-	data = np.load(path+'.npy')
-	with open(path+'.shp') as f: shp = json.load(f)
-	net.load_data(data, shp)
+	net.load_weights(np.load(path+'.npy'))
 	return net
 
 if __name__ == '__main__':
-	pass
+	pass

demo.py

@@ -56,7 +56,7 @@ def forward(self, x):
 shp = {'conv':[[64, 3, 3, 3], [1, 64]]}
 data = np.zeros(64*3*3*3+1*64, dtype=np.float32)
 
-net.load_data(data, shp)
+net.load_weights(data)
 
 load_a_net = '''
 write the layer, flow, shp as json format,