I’ve run into an issue where larger convolutions are slow in pytorch (e.g. a 128x128 convolution with a 512x512 image). Is there something I can do to speed it up, or maybe I’m doing something incorrectly?
import torch
import numpy as np
from skimage.color import rgb2gray
from skimage.data import astronaut
import matplotlib.pyplot as plt
from scipy.signal import fftconvolve
import time

img_now = rgb2gray(astronaut())
psf_extent = 128
x = np.linspace(-1, 1, psf_extent)
y = np.linspace(-1, 1, psf_extent).reshape((-1, 1))
# build asymmetric psf so no specific tricks (e.g. separable, DCT, etc.) can be used
sigma_sqd = (x+2)**5 + (y+2)**5
psf_now = np.exp(-((x+.1)**2 + (y+.1)**2)/(1e-3*sigma_sqd))
tp = torch.Tensor.view(torch.from_numpy(img_now), 1, 1, img_now.shape[0], img_now.shape[1])
tp.requires_grad_(False)
pp = torch.Tensor.view(torch.from_numpy(psf_now), 1, 1, psf_now.shape[0], psf_now.shape[1])
pp.requires_grad_(False)

start = time.time()
blur_img = torch.nn.functional.conv2d(tp, pp, padding=psf_now.shape[0] // 2 - 1)
print('pytorch conv2d elapsed: %fs' % (time.time() - start))
start = time.time()
blur_scipy = fftconvolve(img_now, psf_now, mode='same')
print('scipy fftconvolve elapsed: %fs' % (time.time() - start))

I get  ~25s for pytorch implementation vs ~.04s for scipy, which I could see if pytorch were summing an actual sliding window but it’s my understanding that pytorch (or at least the latest v1.0) infers which type of convolution to do. Even for smaller convolutions, changing psf_exent above, the scipy version seems to still do better. Going to GPU helps, goes to ~2s on my machine, but this doesn’t seem to be the type of operation I should force into GPU.