add checks for input, weight and bias types when using cudnn conv2d (pytorch#1689)

Author: alykhantejani

test/test_nn.py

@@ -1317,6 +1317,35 @@ def test_Conv2d_inconsistent_types(self):
  # but it should work with the same type
  nn.functional.conv2d(inputs.float(), weights.float())
 
+ @unittest.skipIf(not TEST_CUDA, 'CUDA not available')
+ def test_Conv2d_inconsistent_types_on_GPU_without_cudnn(self):
+ inputs = Variable(torch.randn(4, 1, 7, 7).float().cuda())
+ weights = Variable(torch.randn(1, 1, 3, 3).double().cuda())
+ bias = Variable(torch.randn(1).double().cuda())
+
+ torch.backends.cudnn.enabled = False
+ # inconsistent types should raise an exception
+ self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights))
+ self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights.float(), bias))
+
+ # but it should work with the same type
+ nn.functional.conv2d(inputs.float(), weights.float(), bias.float())
+
+ @unittest.skipIf(not TEST_CUDA, 'CUDA not available')
+ @unittest.skipIf(not TEST_CUDNN, 'CUDNN not available')
+ def test_Conv2d_inconsistent_types_on_GPU_with_cudnn(self):
+ inputs = Variable(torch.randn(4, 1, 7, 7).float().cuda())
+ weights = Variable(torch.randn(1, 1, 3, 3).double().cuda())
+ bias = Variable(torch.randn(1).double().cuda())
+
+ torch.backends.cudnn.enabled = True
+ # inconsistent types should raise an exception
+ self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights))
+ self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights.float(), bias))
+
+ # but it should work with the same type
+ nn.functional.conv2d(inputs.float(), weights.float(), bias.float())
+
  def test_Conv2d_missing_argument(self):
  c = nn.Conv2d(3, 3, 3)
  self.assertRaises(RuntimeError, lambda: c(None))

torch/csrc/autograd/functions/convolution.cpp

@@ -1,3 +1,5 @@
+#include <sstream>
+
 #include "convolution.h"
 
 #include "torch/csrc/autograd/variable.h"
@@ -6,6 +8,8 @@
 #include "torch/csrc/nn/THNN_generic.h"
 #include "torch/csrc/utils/auto_gpu.h"
 
+#include "THPP/Type.hpp"
+
 #ifdef WITH_CUDNN
 #include "torch/csrc/cudnn/Conv.h"
 #include "torch/csrc/cudnn/Handles.h"
@@ -120,11 +124,11 @@ static auto view3d(const Tensor& tensor) -> std::unique_ptr<Tensor> {
  return result;
 }
 
+
 auto ConvForward::apply(const variable_list& inputs) -> variable_list {
  check_input_variables("ConvNd", inputs, 3, 2);
  if (is_padding_neg()) throw std::runtime_error("negative padding is not supported");
  if (is_output_padding_neg()) throw std::runtime_error("negative output_padding is not supported");
-
  AutoGPU guard(inputs[0]->data->getDevice());
  auto input = inputs[0]->data->contiguous();
  std::unique_ptr<Tensor> weight(inputs[1]->data->clone_shallow());
@@ -147,6 +151,16 @@ auto ConvForward::apply(const variable_list& inputs) -> variable_list {
 
  if (use_cudnn(*input)) {
 #ifdef WITH_CUDNN
+ if (input->type() != weight->type()){
+ std::stringstream ss;
+ ss << "Input type (" << thpp::toString(input->type()) << ") and weight type (" << thpp::toString(weight->type()) << ") should be the same";
+ throw std::runtime_error(ss.str());
+ }
+ if (bias.get() != NULL && input->type() != bias->type()){
+ std::stringstream ss;
+ ss << "Input type (" << thpp::toString(input->type()) << ") and bias type (" << thpp::toString(bias->type()) << ") should be the same";
+ throw std::runtime_error(ss.str());
+ }
  output = input->newTensor();
  output->resize(output_size(*input, *weight));
  if (transposed) {
@@ -203,7 +217,6 @@ auto ConvBackward::apply(const variable_list& grad_outputs) -> variable_list {
  check_input_variables("ConvNdBackward", grad_outputs, 1);
  if (is_padding_neg()) throw std::runtime_error("negative padding is not supported");
  if (is_output_padding_neg()) throw std::runtime_error("negative output_padding is not supported");
-
  auto input = input_.unpack_data();
  AutoGPU guard(input->getDevice());
  input = input->contiguous();