Gaussian filter of 2D image in Tensorflow 2.0 on GPU

gaussian_filter_tensorflow.py

import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import tensorflow_probability as tfp
from time import time

def toImage(tensor):
    return tensor.reshape(tensor.shape[:2])

def gaussian_kernel(size: int,
                    mean: float,
                    std: float,
                   ):
    """Makes 2D gaussian Kernel for convolution."""
    d = tfp.distributions.Normal(mean, std)
    vals = d.prob(tf.range(start = -size, limit = size + 1, dtype = tf.float32))
    gauss_kernel = tf.einsum('i,j->ij',
                                  vals,
                                  vals)

    return gauss_kernel / tf.reduce_sum(gauss_kernel)

halfx = 5 # This is the "radius" of the kernel, size 2*x+1
std = 0.5

# Make Gaussian Kernel with desired specs.
gauss_kernel = gaussian_kernel(halfx, 0, std)
gauss_kernel = gauss_kernel[:, :, tf.newaxis, tf.newaxis]

img = tf.random.uniform((4096, 4096, 1, 1))
t1 = time()
res = tf.nn.conv2d(img, gauss_kernel, strides=[1, 1, 1, 1], padding="SAME")
img2 = res.numpy()
print(time() - t1)

img = tf.random.uniform((4096, 4096, 1, 1))
t1 = time()
res = tf.nn.conv2d(img, gauss_kernel, strides=[1, 1, 1, 1], padding="SAME")
img2 = res.numpy()
print(time() - t1)


img = tf.random.uniform((4096, 4096, 1, 1))
t1 = time()
res = tf.nn.conv2d(img, gauss_kernel, strides=[1, 1, 1, 1], padding="SAME")
img2 = res.numpy()
print(time() - t1)