For elastic distortion applied to 3D images, the time spent on elastic deformation (which is typically implemented in NumPy, e.g. as in https://gist.github.com/nasimrahaman/8ed04be1088e228c21d51291f47dd1e6) is almost 1min/image. I wonder if there is already an implementation in PyTorch using cuda that speeds things up.
You can just write the Gaussians in terms of convolutions with a Gaussian kernel (i.e. exp(-0.5*(x-center)**2/sigma**2)) times the normalization factor) cut of to some tasteful size (2 sigma or 3 sigma or what).
I played with it some a while ago in 3d. The thing that screws you up is that with larger sigmas (=filter width) in the Gaussian you’ll have rather large convolution kernels and then the usual cuda kernels not terribly efficient. This could be less pronounced in 2d, but for 3d it was kind of slow. There are some easy and some more possibly more elaborate (but possibly also too compute intensive) ways to mitigate that and I wanted to write a blog post on them but I never got around to it.
The other function you need is torch.nn.functional.grid_sample.
Best regards
Thomas