fast maxpool

Author: yunlongdong

infCNN/LayerLib.py

@@ -1,22 +1,20 @@
 from .Core import Layer
-from .conv import img2col
+from .conv import conv
 import numpy as np
 
 class Dense(Layer):
 
     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.zeros((n_out, n_in), dtype='float32')
-        self.bias = np.zeros(n_out, dtype='float32')
+        self.n_in, self.n_out = n_in, n_out
+        self.K =  np.zeros((n_out, n_in), dtype=np.float32)
+        self.bias = np.zeros(n_out, dtype=np.float32)
 
     def forward(self, x):
         self.x = x
-        w = self.K
-        # the dense layer math: y = x*w
-        y = w.dot(x.T)
-        return y.T + self.bias.reshape((1, -1))
+        y = x.dot(self.K.T)
+        y += self.bias.reshape((1,-1))
+        return y
 
     def load_from_torch(self, weights, bias):
         self.K = weights.copy()
@@ -39,15 +37,10 @@ def __init__(self, C_in, C_out, K_s, Stride):
         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 = img2col(X, self.K)
-        out = out.transpose(0, 2, 3, 1)
-        out = out + self.bias
-        return out.transpose(0, 3, 1, 2)
+        out = conv(X, self.K)
+        out += self.bias.reshape((1,-1,1,1))
+        return out
 
     def load_from_torch(self, weights, bias):
         self.K = weights.copy()

infCNN/OpLib.py

@@ -1,6 +1,7 @@
 from .Core import Op 
 import numpy as np
 from .im2col import im2col_indices 
+from .maxpool import maxpool
 
 
 class ReLU(Op):
@@ -86,4 +87,4 @@ def forward(self, x):
 flatten = Flatten()
 sigmoid = Sigmoid()
 softmax = Softmax()
-maxpool = MaxPool()
+maxpool = maxpool 

infCNN/conv.py

@@ -22,10 +22,13 @@ def fill_col(pdimg, idx, colimg):
             s += 1
     return colimg
 
-def img2col(img, core, buf=[np.zeros(1, dtype=np.int32)]):
+def conv(img, core, stride=(1,1), buf=[np.zeros(1, dtype=np.int32)]):
     # new the col_img, if needed
+    strh, strw = stride
     cimg_w = np.cumprod(core.shape[1:])[-1]
-    cimg_h = img.size//img.shape[1]
+    n,c,h,w = img.shape
+    cimg_h = n*(h//strh)*(w//strw)
+    
     if len(buf[0])<cimg_h*cimg_w:
         buf[0] = col_img = np.zeros(cimg_h*cimg_w, dtype=np.int32)
     else:
@@ -34,9 +37,8 @@ def img2col(img, core, buf=[np.zeros(1, dtype=np.int32)]):
 
     # mark where need
     iimg = img.view(dtype=np.int32)
-    # iimg.view(dtype=np.uint8).ravel()[::4]&=0xfe
     iimg &= 0xfffffffe
-    iimg[:,0,:,:] |= 1
+    iimg[:,0,::strh,::strw] |= 1
 
     # ravel the image
     n,c,h,w = np.array(core.shape)
@@ -51,27 +53,22 @@ def img2col(img, core, buf=[np.zeros(1, dtype=np.int32)]):
     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))
+    return rst.reshape((ni, n, hi//strh, wi//strw))
 
 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 &= 0xfffffffe
-    iimg[:,0,:,:] |= 1
-    
+    #img.ravel()[:] = np.arange(3*512*512)
     core = np.zeros((32, 3, 3, 3), dtype=np.float32)
-    core.ravel()[:] = np.arange(3*3*3*32)
-    
+    #core.ravel()[:] = np.arange(3*3*3*32)
+
+    rst1 = conv(img, core, (1,1))
     start = time()
-    rst1 = img2col(img, core)
-    print(time()-start)
+    rst1 = conv(img, core, (1,1))
+    print('jit cost:', time()-start)
 
     start = time()
-    rst = img2col(img, core)
-    print('jit cost: x10', time()-start)
+    rst2 = conv(img, core, (1,1))
+    print('jit cost:', time()-start)

infCNN/conv_s.py

@@ -0,0 +1,80 @@
+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, stride=(1,1), buf=[np.zeros(1, dtype=np.int32)]):
+    # new the col_img, if needed
+    strh, strw = stride
+    cimg_w = np.cumprod(core.shape[1:])[-1]
+    print(img.shape)
+    n,c,h,w = img.shape
+    cimg_h = n*(h//strh)*(w//strw)
+    
+    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 &= 0xfffffffe
+    iimg[:,0,::strh,::strw] |= 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//strh, wi//strw))
+
+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, (1,1))
+    print(time()-start)
+
+    start = time()
+    rst2 = img2col(img, core, (2,2))
+    print('jit cost: x10', time()-start)

infCNN/maxpool.py

@@ -0,0 +1,48 @@
+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
+    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] = max(colimg[s], pdimg[i+j])
+        s += 1
+    return colimg
+
+def maxpool(img, stride=(2,2)):
+    strh, strw = stride
+    n,c,h,w = img.shape
+    cimg_h = n*c*(h//strh)*(w//strw)
+    
+    iimg = img.view(dtype=np.int32)
+    iimg &= 0xfffffffe
+    iimg[:,:,::strh,::strw] |= 1
+    
+    nbs = neighbors(img.shape[1:], (1,)+stride)
+    shp = (n, c, h//strh, w//strw)
+    colimg = np.zeros(shp, dtype=np.int32)
+    fill_col(iimg.ravel(), nbs, colimg.ravel())
+    return colimg.view(np.float32)
+
+if __name__ == '__main__':
+    img = np.zeros((1, 3, 512, 512), dtype=np.float32)
+
+    
+    start = time()
+    rst1 = maxpool(img)
+    print(time()-start)
+
+    start = time()
+    rst = maxpool(img)
+    print('jit cost: x10', time()-start)