Repository: pytorch/pytorch

Issue:
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


Discussion:
The 1024 limitation was removed in 5.1.10

Resolving patch:
diff --git a/torch/nn/_functions/batchnorm.py b/torch/nn/_functions/batchnorm.py
deleted file mode 100644
--- a/torch/nn/_functions/batchnorm.py
+++ /dev/null
@@ -1,78 +0,0 @@
-import torch
-from torch.autograd.function import Function
-from torch._thnn import type2backend
-import torch.backends.cudnn as cudnn
-
-
-class BatchNorm(Function):
-
-    def __init__(self, running_mean, running_var, training, momentum, eps):
-        super(BatchNorm, self).__init__()
-        self.running_mean = running_mean
-        self.running_var = running_var
-        self.training = training
-        self.momentum = momentum
-        self.eps = eps
-
-    def forward(self, input, weight=None, bias=None):
-        self.save_for_backward(input, weight, bias)
-
-        # 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
-                          weight is not None and bias is not None and
-                          not isinstance(input, torch.cuda.HalfTensor))
-
-        # temporary buffers used in forward and backward
-        num_features = input.size(1)
-        _save_mean = input.new(num_features)
-        _save_std = input.new(num_features)
-
-        output = input.new(input.size())
-
-        if self.use_cudnn:
-            torch._C._cudnn_batch_norm_forward(
-                input, output, weight, bias,
-                self.running_mean, self.running_var, _save_mean,
-                _save_std, self.training, self.momentum, self.eps)
-        else:
-            backend = type2backend[type(input)]
-            backend.BatchNormalization_updateOutput(
-                backend.library_state, input, output, weight, bias,
-                self.running_mean, self.running_var, _save_mean,
-                _save_std, self.training, self.momentum, self.eps)
-
-        if self.requires_grad:
-            self._save_mean = _save_mean
-            self._save_std = _save_std
-
-        return output
-
-    def backward(self, grad_output):
-        input, weight, bias = self.saved_tensors
-        grad_input, grad_weight, grad_bias = None, None, None
-
-        if self.needs_input_grad[0] or self.use_cudnn:
-            grad_input = input.new(input.size())
-        if (len(self.needs_input_grad) > 1 and self.needs_input_grad[1]) or self.use_cudnn:
-            grad_weight = weight.new(weight.size()).zero_()
-        if (len(self.needs_input_grad) > 1 and self.needs_input_grad[2]) or self.use_cudnn:
-            grad_bias = bias.new(bias.size()).zero_()
-
-        if self.use_cudnn and self.training:
-            # cudnn does not support backward in evaluate mode
-            torch._C._cudnn_batch_norm_backward(
-                input, grad_output, grad_input,
-                grad_weight, grad_bias, weight,
-                self.running_mean, self.running_var,
-                self._save_mean, self._save_std, self.training, self.eps)
-        else:
-            grad_output = grad_output.contiguous()
-            backend = type2backend[type(input)]
-            backend.BatchNormalization_backward(
-                backend.library_state, input, grad_output, grad_input,
-                grad_weight, grad_bias, weight,
-                self.running_mean, self.running_var,
-                self._save_mean, self._save_std, self.training, 1.0, self.eps)
-
-        return grad_input, grad_weight, grad_bias
diff --git a/torch/nn/backends/thnn.py b/torch/nn/backends/thnn.py
--- a/torch/nn/backends/thnn.py
+++ b/torch/nn/backends/thnn.py
@@ -21,7 +21,6 @@ def _get_thnn_function_backend():
 def _initialize_backend():
     from .._functions.thnn import _all_functions as _thnn_functions
     from .._functions.linear import Linear
-    from .._functions.batchnorm import BatchNorm
     from .._functions.conv import ConvNd
     from .._functions.rnn import RNN, \
         RNNTanhCell, RNNReLUCell, GRUCell, LSTMCell
@@ -31,7 +30,6 @@ def _initialize_backend():
         HingeEmbeddingLoss, MarginRankingLoss
 
     backend.register_function('Linear', Linear)
-    backend.register_function('BatchNorm', BatchNorm)
     backend.register_function('ConvNd', ConvNd)
     backend.register_function('RNN', RNN)
     backend.register_function('RNNTanhCell', RNNTanhCell)


Test patch: