adding upsample support

Author: yunlongdong

infCNN/LayerLib.py

@@ -1,28 +1,29 @@
 from .Core import Layer
-from .conv import conv
+from .conv_s 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, self.n_out = n_in, n_out
-        self.K =  np.zeros((n_out, n_in), dtype=np.float32)
+        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
         y = x.dot(self.K.T)
-        y += self.bias.reshape((1,-1))
+        y += self.bias.reshape((1, -1))
         return y
 
     def load_from_torch(self, weights, bias):
         self.K = weights.copy()
         self.bias = bias.copy()
-    
+
 
 class Conv2d(Layer):
-    
+
     def __init__(self, C_in, C_out, K_s, Stride):
         """
         Params:
@@ -33,27 +34,24 @@ def __init__(self, C_in, C_out, K_s, Stride):
         self.c_out = C_out
         self.k_s = K_s
         self.stride = Stride
-        self.pad_size = int((K_s -1)/2)
+        self.pad_size = int((K_s - 1)/2)
         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(X, self.K)
-        out += self.bias.reshape((1,-1,1,1))
+        out = conv(X, self.K, (self.stride, self.stride))
+        out += self.bias.reshape((1, -1, 1, 1))
         return out
-        
+
     def load_from_torch(self, weights, bias):
         self.K = weights.copy()
         self.bias = bias.copy()
 
 
-
-
 if __name__ == "__main__":
     # data = np.arange(16).reshape((1, 1, 4, 4))
     data = np.random.randn(2, 3, 200, 200)
 
-
     import torch.nn as nn
     import torch
     from time import time
@@ -68,7 +66,7 @@ def load_from_torch(self, weights, bias):
     start = time()
     out_torch = conv_torch(torch.FloatTensor(data)).detach().numpy()
     print('torch time:', time() - start)
-    
+
     conv.load_from_torch(w, b)
     start = time()
     out = conv(data)
@@ -78,7 +76,6 @@ def load_from_torch(self, weights, bias):
     print('acw:', out.shape)
     print(np.sum(out-out_torch))
 
-
     data = np.random.randn(2, 10)
 
     dense_torch = nn.Linear(10, 20)

infCNN/OpLib.py

@@ -1,17 +1,17 @@
-from .Core import Op 
+from .Core import Op
 import numpy as np
-from .im2col import im2col_indices 
 from .maxpool import maxpool
+from .upsample import upsample
+
 
-    
 class ReLU(Op):
 
     def __init__(self):
         super().__init__('ReLU')
 
     def forward(self, x):
         self.x = x
-        return (x>0) * x
+        return (x > 0) * x
 
 
 class Flatten(Op):
@@ -34,6 +34,7 @@ def forward(self, x):
         self.x = x
         return 1/(1 + np.exp(-x))
 
+
 class Softmax(Op):
     def __init__(self):
         super().__init__('Sofmax')
@@ -43,48 +44,9 @@ def forward(self, X, axis=-1):
         return (eX.T / eX.sum(axis=axis)).T
 
 
-
-
-class MaxPool(Op):
-
-    def __init__(self, size=2, stride=2):
-        super().__init__('MaxPool_{}x{}s{}'.format(size, size, stride))
-        self.size = size
-        self.stride = stride
-
-    def _pool_forward(self, X, pool_fun, size=2, stride=2):
-        n, d, h, w = X.shape
-        h_out = (h - size) / stride + 1
-        w_out = (w - size) / stride + 1
-    
-        if not w_out.is_integer() or not h_out.is_integer():
-            raise Exception('Invalid output dimension!')
-    
-        h_out, w_out = int(h_out), int(w_out)
-    
-        X_reshaped = X.reshape(n * d, 1, h, w)
-        X_col = im2col_indices(X_reshaped, size, size, padding=0, stride=stride)
-    
-        out, pool_cache = pool_fun(X_col)
-    
-        out = out.reshape(h_out, w_out, n, d)
-        out = out.transpose(2, 3, 0, 1)
-    
-        cache = (X, size, stride, X_col, pool_cache)
-    
-        return out# , cache
-    
-    def maxpool(self, X_col):
-        max_idx = np.argmax(X_col, axis=0)
-        out = X_col[max_idx, range(max_idx.size)]
-        return out, max_idx
-
-    def forward(self, x):
-        return self._pool_forward(x, self.maxpool, self.size, self.stride)
-
-
 relu = ReLU()
 flatten = Flatten()
 sigmoid = Sigmoid()
 softmax = Softmax()
-maxpool = maxpool 
+maxpool = maxpool
+upsample = upsample

infCNN/__init__.py

@@ -1,3 +1,2 @@
 from .LayerLib import * 
 from .OpLib import * 
-from .numpy_extend import *

infCNN/conv_s.py

@@ -2,61 +2,66 @@
 from time import time
 from numba import njit
 
+
 def neighbors(shape, core):
-    shp = [slice(0,i) for i in core]
+    shp = [slice(0, i) for i in core]
     idx = np.mgrid[tuple(shp)]
-    idx = idx.reshape((len(core),-1))
+    idx = idx.reshape((len(core), -1))
     offset = np.array(core)//2
     offset[0] = 0
-    idx -= offset.reshape((-1,1))
+    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
+        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)]):
+
+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]
-    print(img.shape)
-    n,c,h,w = img.shape
+    # print(img.shape)
+    n, c, h, w = img.shape
     cimg_h = n*(h//strh)*(w//strw)
-    
-    if len(buf[0])<cimg_h*cimg_w:
+
+    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.view(dtype=np.uint8).ravel()[::4]&=0xfe
     iimg &= 0xfffffffe
-    iimg[:,0,::strh,::strw] |= 1
-    
+    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))
+    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))
+    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
@@ -65,16 +70,16 @@ def img2col(img, core, stride=(1,1), buf=[np.zeros(1, dtype=np.int32)]):
     img.ravel()[:] = np.arange(3*512*512)
 
     iimg = img.view(dtype=np.int32)
-    iimg.view(dtype=np.uint8).ravel()[::4]&=0xfe
-    iimg[:,0,:,:] |= 1
-    
+    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))
+    rst1 = img2col(img, core, (1, 1))
     print(time()-start)
 
     start = time()
-    rst2 = img2col(img, core, (2,2))
+    rst2 = img2col(img, core, (2, 2))
     print('jit cost: x10', time()-start)

infCNN/upsample.py

@@ -0,0 +1,50 @@
+import numpy as np
+from numba import njit
+from math import floor
+
+@njit
+def _resize(img, k, ra, rs, _rs, ca, cs, _cs, out):
+    h, w = img.shape
+    for r in range(h*k):
+        rar = ra[r]
+        rbr = rar + 1
+        rsr = rs[r]
+        _rsr = _rs[r]
+        for c in range(w*k):
+            cac = ca[c]
+            cbc = cac + 1
+            rra = img[rar,cac]*_rsr
+            rra += img[rbr,cac]*rsr
+            rrb = img[rar,cbc]*_rsr
+            rrb += img[rbr,cbc]*rsr
+            rcab = rra * _cs[c] + rrb * cs[c]
+            out[r,c] = rcab
+
+def upsample(img, k=2, out=None):
+    nc, (h, w), e = img.shape[:-2], img.shape[-2:], 1e-4
+    if out is None:
+        out = np.zeros(nc+(h*k, w*k), dtype=img.dtype)
+    
+    rs = np.linspace(e,h-1-e, h*k, dtype=np.float32)
+    cs = np.linspace(e,w-1-e, w*k, dtype=np.float32)
+    ra = np.floor(rs).astype(np.uint32)
+    ca = np.floor(cs).astype(np.uint32)
+    rs -= ra
+    cs -= ca
+
+    outcol = out.reshape((-1, h*k, w*k))
+    imgcol = img.reshape((-1, h, w))
+    for i, o in zip(imgcol, outcol):
+        _resize(i, k, ra, rs, 1-rs, ca, cs, 1-cs, o)
+    return out
+
+if __name__ == '__main__':
+    from time import time
+    from skimage.data import astronaut
+    img = astronaut().transpose(2,0,1).astype(np.float32)
+
+    # firs time to jit
+    resize(img, 2)
+    start = time()
+    rst = resize(img, 2)
+    print('\nmine v1', time()-start)