numpy conv2d tricks

conv2dtricks.py

import numpy as np
from numpy.lib.stride_tricks import as_strided
import tensorflow as tf
import time

def conv2dTrickster(a, b):
    a = as_strided(a,(len(a),a.shape[1]-len(b)+1,a.shape[2]-b.shape[1]+1,len(b),b.shape[1],a.shape[3]),a.strides[:3]+a.strides[1:])
    return np.einsum('abcijk,ijkd', a, b[::-1,::-1])
    
def conv2dSimple(image, filter):
    # Height and width of output image
    Hout = image.shape[1] - filter.shape[0] + 1
    Wout = image.shape[2] - filter.shape[1] + 1

    output = np.zeros([image.shape[0], Hout, Wout, filter.shape[3]])

    for n in range(output.shape[0]):
        for i in range(output.shape[1]):
            for j in range(output.shape[2]):
                for cout in range(output.shape[3]):
                    output[n,i,j,cout] = np.multiply(image[n, i:i+filter.shape[0], j:j+filter.shape[1], :], filter[:,:,:,cout]).sum()

    return output
  
def tfconv2d(a, b):
    out = tf.nn.conv2d(a, b, strides=[1,1,1,1], padding="VALID")
    return out

imagesizes = [64, 128, 256, 512, 1024]
for imagesize in imagesizes:
    image  = np.random.ranf([imagesize, imagesize, 3, 3])
    filter = np.random.ranf([3,   3,   3, 3])
    # Tensorflow
    startTime = time.time()
    tfout = tfconv2d(image, filter)
    with tf.Session() as sess:
        sess.run(tfout)
        print ("Tensorflow time: [" + str(imagesize) + "]" + str(time.time() - startTime) + ", shape = " + str(tfout.get_shape()) + "tfout =" + str(tfout.eval()[0][0]))
    # stride trick
    startTime = time.time()
    trickout = conv2dTrickster(image, filter)
    print ("Trick time: : [" + str(imagesize) + "]" + str(time.time() - startTime) + ", trickout = " + str(trickout[0][0]))
    # regular trick
    startTime = time.time()
    simpleout = conv2dSimple(image, filter)
    print ("Simple time: : [" + str(imagesize) + "]" + str(time.time() - startTime) + ", simpleout = " + str(simpleout[0][0]))

"""
2019-05-14 10:37:34.715008: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2
Tensorflow time: [64]0.13351082801818848, shape = (64, 62, 1, 3)tfout =[[7.30340124 7.44151118 7.19464267]]
Trick time: : [64]0.027161121368408203, trickout = [[7.44536471 7.09516103 7.29013894]]
Simple time: : [64]0.251690149307251, simpleout = [[7.30340124 7.44151118 7.19464267]]
Tensorflow time: [128]0.12033915519714355, shape = (128, 126, 1, 3)tfout =[[7.90427422 9.01649472 9.03089107]]
Trick time: : [128]0.10680913925170898, trickout = [[7.44665604 7.87937458 9.37562679]]
Simple time: : [128]1.0285594463348389, simpleout = [[7.90427422 9.01649472 9.03089107]]
Tensorflow time: [256]0.4307737350463867, shape = (256, 254, 1, 3)tfout =[[6.39663357 5.9422466  6.6824048 ]]
Trick time: : [256]0.4327685832977295, trickout = [[6.56362405 5.89000768 7.07348136]]
Simple time: : [256]4.267130613327026, simpleout = [[6.39663357 5.9422466  6.6824048 ]]
Tensorflow time: [512]1.7026965618133545, shape = (512, 510, 1, 3)tfout =[[5.22365016 6.18046664 5.76867063]]
Trick time: : [512]1.7334675788879395, trickout = [[6.06602827 6.15500158 5.83996214]]
Simple time: : [512]16.582613468170166, simpleout = [[5.22365016 6.18046664 5.76867063]]
Tensorflow time: [1024]12.834256887435913, shape = (1024, 1022, 1, 3)tfout =[[6.23548267 5.60928775 4.30881343]]
Trick time: : [1024]7.12175178527832, trickout = [[5.77306404 5.11819285 3.89802783]]
Simple time: : [1024]67.58800268173218, simpleout = [[6.23548267 5.60928775 4.30881343]]
"""

np.tensordot(a, b, axes=3) is faster