Where is torch.cudnn_affine_grid_generator?

I want to modify the affine_grid_generator function in torch/nn/_functions/vision.py from 2D to 3D. But I can’t find where is torch.cudnn_affine_grid_generator function. My pytorch version is 0.4.0 and python version is 3.6. The source code is shown below. Thanks in advance!

def affine_grid_generator(theta, size):
if theta.data.is_cuda:
if not cudnn.enabled:
raise RuntimeError("AffineGridGenerator needs CuDNN for "
“processing CUDA inputs, but CuDNN is not enabled”)
if not cudnn.is_acceptable(theta.data):
raise RuntimeError(“AffineGridGenerator generator theta not acceptable for CuDNN”)
N, C, H, W = size
return torch.cudnn_affine_grid_generator(theta, N, C, H, W)
return AffineGridGenerator.apply(theta, size)

def affine_grid_generator_3d(theta, size):
return AffineGridGenerator_3d.apply(theta, size)

It’s in the C++ modules (ATen to be precise).

Personally, I tend to use the “rgrep” command on the PyTorch source tree to find functions.

Best regards


Hi @Keke
Did you solve this issue. I want to do exactly the same.

I am also interested in a solution to this. @Keke did you happen to build a 3D affine grid generator?

I have the same question! How did you do!

This looks like it points to cudnn code, so I can’t access the source. Does anyone know where the actual code is stored, maybe for the cpu version? I want to see the code itself for my experiments.

Yeah, well. The question was about cudnn_affine_grid_generator.

Without CuDNN, by applying the same method to search, it turns out that /aten/src/ATen/native/AffineGridGenerator.cpp is the code that is used, independent of CPU vs. CUDA.

1 Like