faster cpu conv

Author: yunlongdong

infCNN/LayerLib.py

@@ -1,5 +1,5 @@
 from .Core import Layer
-from .im2col import im2col_indices
+from .conv import img2col
 import numpy as np
 
 class Dense(Layer):
@@ -8,8 +8,8 @@ def __init__(self, n_in, n_out):
         super().__init__('dense_{}_{}'.format(n_in, n_out))
         self.n_in = n_in
         self.n_out = n_out
-        self.K =  np.random.randn(n_out, n_in)
-        self.bias = np.zeros(n_out)
+        self.K =  np.zeros((n_out, n_in), dtype='float32')
+        self.bias = np.zeros(n_out, dtype='float32')
 
     def forward(self, x):
         self.x = x
@@ -30,50 +30,28 @@ def __init__(self, C_in, C_out, K_s, Stride):
         Params:
             in channels, out channles, kernel size, stride
         """
-        super().__init__("conv_{}_{}x{}x{}".format(C_out, C_in, K_s, K_s))
+        super().__init__("conv_{}_{}x{}x{}".format(C_in, C_out, K_s, K_s))
         self.c_in = C_in
         self.c_out = C_out
         self.k_s = K_s
         self.stride = Stride
         self.pad_size = int((K_s -1)/2)
-        self.K = np.zeros((C_out, C_in, K_s, K_s))
-        self.bias = np.zeros(C_out)
+        self.K = np.zeros((C_out, C_in, K_s, K_s), dtype='float32')
+        self.bias = np.zeros(C_out, dtype='float32')
 
-    def forward(self, X):
-        out = conv_forward(X, self.K, self.bias, stride=self.stride, padding=self.pad_size)
-        return out
+    # def forward(self, X):
+    #     out = conv_forward(X, self.K, self.bias, stride=self.stride, padding=self.pad_size)
+    #     return out
 
+    def forward(self, X):
+        out = img2col(X, self.K)
+        print(out.shape)
 
     def load_from_torch(self, weights, bias):
         self.K = weights
         self.bias = bias
 
 
-def conv_forward(X, W, b, stride=1, padding=1):
-    cache = W, b, stride, padding
-    n_filters, d_filter, h_filter, w_filter = W.shape
-    n_x, d_x, h_x, w_x = X.shape
-    h_out = (h_x - h_filter + 2 * padding) / stride + 1
-    w_out = (w_x - w_filter + 2 * padding) / stride + 1
-
-    if not h_out.is_integer() or not w_out.is_integer():
-        raise Exception('Invalid output dimension!')
-
-    h_out, w_out = int(h_out), int(w_out)
-
-    X_col = im2col_indices(X, h_filter, w_filter, padding=padding, stride=stride)
-    W_col = W.reshape(n_filters, -1)
-
-    out = W_col @ X_col
-    out = (out.T + b[:]).T
-    out = out.reshape(n_filters, h_out, w_out, n_x)
-    out = out.transpose(3, 0, 1, 2)
-
-    cache = (X, W, b, stride, padding, X_col)
-
-    return out# , cache
-
-
 
 
 if __name__ == "__main__":

infCNN/conv.py

@@ -0,0 +1,75 @@
+import numpy as np
+from time import time
+from numba import njit
+
+def neighbors(shape, core):
+    shp = [slice(0,i) for i in core]
+    idx = np.mgrid[tuple(shp)]
+    idx = idx.reshape((len(core),-1))
+    offset = np.array(core)//2
+    offset[0] = 0
+    idx -= offset.reshape((-1,1))
+    acc = np.cumprod((1,)+shape[::-1][:-1])
+    return np.dot(idx.T, acc[::-1])
+
+@njit
+def fill_col(pdimg, idx, colimg):
+    s = 0
+    for i in range(len(pdimg)):
+        if pdimg[i]&1==0: continue
+        for j in idx:
+            colimg[s] = pdimg[i+j]
+            s += 1
+    return colimg
+
+def img2col(img, core, buf=[np.zeros(1, dtype=np.int32)]):
+    # new the col_img, if needed
+    cimg_w = np.cumprod(core.shape[1:])[-1]
+    cimg_h = img.size//img.shape[1]
+    if len(buf[0])<cimg_h*cimg_w:
+        buf[0] = col_img = np.zeros(cimg_h*cimg_w, dtype=np.int32)
+    else:
+        col_img = buf[0][:cimg_h*cimg_w]
+        col_img[:] = 0
+    
+    # mark where need
+    iimg = img.view(dtype=np.int32)
+    iimg.view(dtype=np.uint8).ravel()[::4]&=0xfe
+    iimg[:,0,:,:] |= 1
+    
+    # ravel the image
+    n,c,h,w = np.array(core.shape)
+    shp = ((0,0),(0,0),(h//2,h//2),(w//2,w//2))
+    pdimg = np.pad(iimg, shp, 'constant', constant_values=0)
+    nbs = neighbors(pdimg.shape[1:], core.shape[1:])
+    fill_col(pdimg.ravel(), nbs, col_img)
+    col_img = col_img.view(np.float32)
+    col_img = col_img.reshape((cimg_h, cimg_w))
+    
+    # dot
+    col_core = core.reshape((core.shape[0],-1))
+    rst = col_core.dot(col_img.T)
+    ni, ci, hi, wi = img.shape
+    return rst.reshape((ni, n, hi, wi))
+
+if __name__ == '__main__':
+    from skimage.data import camera
+    import matplotlib.pyplot as plt
+    from scipy.ndimage import convolve
+    img = np.zeros((1, 3, 512, 512), dtype=np.float32)
+    img.ravel()[:] = np.arange(3*512*512)
+
+    iimg = img.view(dtype=np.int32)
+    iimg.view(dtype=np.uint8).ravel()[::4]&=0xfe
+    iimg[:,0,:,:] |= 1
+    
+    core = np.zeros((32, 3, 3, 3), dtype=np.float32)
+    core.ravel()[:] = np.arange(3*3*3*32)
+    
+    start = time()
+    rst1 = img2col(img, core)
+    print(time()-start)
+
+    start = time()
+    rst = img2col(img, core)
+    print('jit cost: x10', time()-start)

infCNN/im2col.py

@@ -18,6 +18,14 @@ def up4x(fm):
     fm4x = up2x(up2x(fm))
     return fm4x
 
+def up8x(fm):
+    fm8x = up4x(up2x(fm))
+    return fm8x
+
+def up16x(fm):
+    fm16x = up8x(up2x(fm))
+    return fm16x
+
 if __name__ == "__main__":
     img = np.random.randn(2, 3, 100, 100)
     img2x = up4x(img)