Repository: pytorch/pytorch
Pull request: dont use cudnn batchnorm for cudnn < 5.1.10
fixes #1188 

Problem statement:
cudnn < 6 needs to disable batchnorm codepath
As Arthur Szlam reports, fb-internal cudnn is still lagging behind, and giving batch-size > 1024 with batchnorm is raising an error from cudnn. Need to check for compile-time version and disable this codepath

Hints from issue discussion:
The 1024 limitation was removed in 5.1.10

Commit message: dont use cudnn batchnorm for cudnn < 5.1.10

Patch to review:
diff --git a/torch/csrc/autograd/functions/batch_normalization.cpp b/torch/csrc/autograd/functions/batch_normalization.cpp
index 0722807a1b97..20cc47371c5e 100644
--- a/torch/csrc/autograd/functions/batch_normalization.cpp
+++ b/torch/csrc/autograd/functions/batch_normalization.cpp
@@ -27,7 +27,8 @@ auto BatchNormForward::apply(const variable_list& inputs) -> variable_list {
 #ifdef WITH_CUDNN
   use_cudnn = (input->data->isCuda()
                && input->data->type() != thpp::Type::HALF
-               && weight && bias && cudnn_enabled);
+               && weight && bias
+               && cudnn_enabled && CUDNN_VERSION >= 5110L);
 #endif
 
   auto output = input->data->newTensor();
@@ -91,7 +92,8 @@ auto BatchNormBackward::apply(const variable_list& grad_outputs) -> variable_lis
 #ifdef WITH_CUDNN
   use_cudnn = (input->isCuda()
                && input->type() != thpp::Type::HALF
-               && weight && bias && training && cudnn_enabled);
+               && weight && bias && training
+               && cudnn_enabled && CUDNN_VERSION >= 5110L);
 #endif
 
   std::unique_ptr<Tensor> grad_input;
diff --git a/torch/nn/_functions/batchnorm.py b/torch/nn/_functions/batchnorm.py
index db8f7a036482..e949d234d044 100644
--- a/torch/nn/_functions/batchnorm.py
+++ b/torch/nn/_functions/batchnorm.py
@@ -20,6 +20,7 @@ def forward(self, input, weight=None, bias=None):
         # don't use cuDNN for half inputs because cuDNN requires the weight and
         # bias tensors to be floats, unlike THCUNN which requires half tensors.
         self.use_cudnn = (cudnn.is_acceptable(input) and
+                          cudnn.version() > 5110 and
                           weight is not None and bias is not None and
                           not isinstance(input, torch.cuda.HalfTensor))
 


Code review comments:
diff_hunk: @@ -20,6 +20,7 @@ def forward(self, input, weight=None, bias=None):
         # don't use cuDNN for half inputs because cuDNN requires the weight and
         # bias tensors to be floats, unlike THCUNN which requires half tensors.
         self.use_cudnn = (cudnn.is_acceptable(input) and
+                          cudnn.version() > 5110 and
                           weight is not None and bias is not None and
                           not isinstance(input, torch.cuda.HalfTensor))
path: torch/nn/_functions/batchnorm.py
text: @user1:
Huh, it seems weird that `BatchNormBackward` checks the `training` flag, but the Python function doesn't. Is there any problem with using cuDNN at inference?

Also, we should probably get rid of this and only use the C one. 

@author:
we cannot use cudnn for backward in inference mode, as it computes incorrectly the gradInput in inference mode

@user1:
but is it ok to use cuDNN for forward and THCUNN for backward?

@author:
oh hmm, sam says probably not. i'll make them consistent.

@user3:
Why? THCUNN with affine transformation and cuDNN should be fully compatible. @user2 had PR #901 that would make batchNorm for inference unnecessary, everyone should use that, it's faster.