Deconvolutions improve performance of AH-Net

[mmarcinkiewicz]

Describe the bug
In the current environment deconvolutions work significantly faster with mixed precision than trilinear interpolation for upsampling. In the current Monai implementation of AH-Net the upsample_mode defaults to "trilinear", which is suboptimal.

To Reproduce
Run training with upsample_mode="trilinear" and upsample_mode="transpose".

Expected behavior
The default upsample_mode is changed to transpose.

Screenshots
Throughput measured on V100 and A100 for a 3D workload.

Environment (please complete the following information):

• Ubuntu 18.04
• Python 3.6
• MONAI version f1998b72a941d1e5f9578a66dc1c20b01913caab
• CUDA 11, cuDNN 8.0.3/8.0.4
• GPUs: V100, A100

Additional context
In cuDNN 8.0.3 deconvolutions have Tensor Core support which improves their performance significantly.
Currently there are issues being tracked for Pytorch and cuDNN which might improve performance. I'll monitor the status and update if needed.

[Nic-Ma]

Hi @yiheng-wang-nv ,

I remember you fixed some determinism issue about the upsampling in AHNet, is it related to this ticket?

Thanks.

[wyli]

Hi @yiheng-wang-nv ,

I remember you fixed some determinism issue about the upsampling in AHNet, is it related to this ticket?

Thanks.

that's noted in the docstring @Nic-Ma

MONAI/monai/networks/nets/ahnet.py

Lines 353 to 355 in 8fc4f5d

	upsample_mode: [``"transpose"``, ``"bilinear"``, ``"trilinear"``]
	The mode of upsampling manipulations.
	Using the last two modes cannot guarantee the model's reproducibility. Defaults to ``trilinear``.

[yiheng-wang-nv]

Hi @yiheng-wang-nv ,

I remember you fixed some determinism issue about the upsampling in AHNet, is it related to this ticket?

Thanks.

FYI, the non-determinism issue is explained in an old PR. Generally speaking, torch.nn.functional.interpolate() and torch.nn.MaxPool1d/2d/3d() (if kernel size != stride) use atomicAdd which introduces the non-determinism results.
This ticket is about the comparisons for computing speed between trilinear and deconvolution, and since our deconvolution mode ensures the determinism, the kernel size for maxpool layer has been changed (shown below), will it have an impact on the speed?

if upsample_mode == "transpose":
    self.maxpool = pool_type(kernel_size=(2, 2, 2)[-spatial_dims:], stride=2)
else:
    self.maxpool = pool_type(kernel_size=(3, 3, 3)[-spatial_dims:], stride=2, padding=1)