addressing comments

Author: adamlerer

test/test_autograd.py

@@ -230,6 +230,20 @@ def test_setitem_mask(self):
  mask = Variable(torch.ByteTensor(5, 5).bernoulli_(), requires_grad=False)
  self._test_setitem(mask)
 
+ def test_setitem_tensor(self):
+ x = Variable(torch.ones(5, 5), requires_grad=True)
+ y = x + 2
+ y_version = y._version
+ value = Variable(torch.Tensor(5).fill_(7), requires_grad=True)
+ index = 3
+ y[index] = value
+ self.assertNotEqual(y._version, y_version)
+ y.backward(torch.ones(5, 5))
+ expected_grad_input = torch.ones(5, 5)
+ expected_grad_input[index] = 0
+ self.assertEqual(x.grad, expected_grad_input)
+ self.assertEqual(value.grad, torch.ones(5))
+
  def test_type_conversions(self):
  import torch.cuda
  x = Variable(torch.randn(5, 5))

test/test_nn.py

@@ -5,6 +5,7 @@
 import contextlib
 from copy import deepcopy
 from itertools import repeat
+from functools import wraps
 
 import torch.nn as nn
 import torch.nn.parallel as dp
@@ -13,14 +14,19 @@
  module_tests, criterion_tests, TEST_CUDA, TEST_CUDNN, PRECISION
 from common import freeze_rng_state
 
-@contextlib.contextmanager
-def set_default_tensor_type(type):
- old_type = torch.typename(torch.Tensor())
- torch.set_default_tensor_type(type)
- try:
- yield
- finally:
- torch.set_default_tensor_type(old_type)
+def default_tensor_type(type):
+ type_str = torch.typename(type)
+ def decorator(fn):
+ @wraps(fn)
+ def wrapper(*args, **kwargs):
+ old_type = torch.typename(torch.Tensor())
+ torch.set_default_tensor_type(type_str)
+ try:
+ return fn(*args, **kwargs)
+ finally:
+ torch.set_default_tensor_type(old_type)
+ return wrapper
+ return decorator
 
 class InputVariableMixin(object):
  def _get_input(self):
@@ -621,7 +627,7 @@ def test_MaxUnpool2d_output_size(self):
  def test_RNN_cell(self):
  # this is just a smoke test; these modules are implemented through
  # autograd so no Jacobian test is needed
- for module in (nn.rnn.cell.RNN, nn.rnn.cell.RNNReLU, nn.rnn.cell.GRU):
+ for module in (nn.RNNCell, nn.GRUCell):
  for bias in (True, False):
  input = Variable(torch.randn(3, 10))
  hx = Variable(torch.randn(3, 20))
@@ -638,18 +644,18 @@ def test_LSTM_cell(self):
  input = Variable(torch.randn(3, 10))
  hx = Variable(torch.randn(3, 20))
  cx = Variable(torch.randn(3, 20))
- lstm = nn.rnn.cell.LSTM(10, 20, bias=bias)
+ lstm = nn.LSTMCell(10, 20, bias=bias)
  for i in range(6):
  hx, cx = lstm(input, (hx, cx))
 
  (hx+cx).sum().backward()
 
  @unittest.skipIf(not TEST_CUDNN, "needs cudnn")
+ @default_tensor_type(torch.FloatTensor) # FIXME: just until torch.cuda.DoubleTensor.sum() implemented
  def test_RNN_cpu_vs_cudnn(self):
 
- def forward_backward(cuda, module, bias, input_val, hx_val, weights_val):
- rnn = module(input_size, hidden_size, num_layers, bias=bias)
- is_lstm = module == nn.rnn.LSTM
+ def forward_backward(cuda, rnn, input_val, hx_val, weights_val):
+ is_lstm = type(rnn) == nn.LSTM
 
  for x_layer, y_layer in zip(rnn.all_weights, weights_val):
  for x, y in zip(x_layer, y_layer):
@@ -689,42 +695,46 @@ def forward_backward(cuda, module, bias, input_val, hx_val, weights_val):
  'cy': hy[1].data if is_lstm else None,
  'grad_cx': hx[1].grad if is_lstm else None}
 
- def diff(t_cpu, t_gpu, name):
- self.assertTrue(torch.is_tensor(t_cpu))
- self.assertTrue(torch.is_tensor(t_gpu))
- delta = t_gpu.cpu().add(-1, t_cpu).abs().max()
- # print("{:30s} cpu: {:10g} gpu: {:10g} diff: {:10g}".format(name, t_cpu.abs().max(), t_gpu.abs().max(), delta))
- self.assertLess(delta, 2 * PRECISION)
-
  input_size = 10
  hidden_size = 20
  num_layers = 2
  seq_length = 7
  batch = 5
 
- # FIXME: we can't use torch.cuda.DoubleTensor because sum() is not yet defined on it
- with set_default_tensor_type('torch.FloatTensor'):
- for module in (nn.rnn.RNNTanh, nn.rnn.RNNReLU, nn.rnn.LSTM, nn.rnn.GRU):
- for bias in (True, False):
- input_val = torch.randn(seq_length, batch, input_size)
- hx_val = torch.randn(num_layers, batch, hidden_size)
-
- weights_val = module(input_size, hidden_size, num_layers).all_weights
-
- outputs_cpu = forward_backward(False, module, bias, input_val, hx_val, weights_val)
- outputs_gpu = forward_backward(True, module, bias, input_val, hx_val, weights_val)
-
- diff(outputs_cpu['output'], outputs_gpu['output'], 'output')
- diff(outputs_cpu['hy'], outputs_gpu['hy'], 'hy')
- diff(outputs_cpu['grad_input'], outputs_gpu['grad_input'], 'grad_input')
- diff(outputs_cpu['grad_hx'], outputs_gpu['grad_hx'], 'grad_hx')
- if outputs_cpu['cy'] is not None:
- diff(outputs_cpu['cy'], outputs_gpu['cy'], 'cy')
- diff(outputs_cpu['grad_cx'], outputs_gpu['grad_cx'], 'grad_cx')
-
- for i, (cpu_layer_weight, gpu_layer_weight) in enumerate(zip(outputs_cpu['weights'], outputs_gpu['weights'])):
- for j, (cpu_weight, gpu_weight) in enumerate(zip(cpu_layer_weight, gpu_layer_weight)):
- diff(cpu_weight.grad, gpu_weight.grad, 'grad_weight[{},{}]'.format(i, j))
+ def compare_cpu_gpu(outputs_cpu, outputs_gpu):
+ self.assertEqual(list(outputs_cpu.keys()), list(outputs_gpu.keys()))
+ for key in outputs_cpu.keys():
+ if key != 'weights':
+ self.assertEqual(outputs_cpu[key], outputs_gpu[key], prec=5e-5)
+
+ # check grad weights separately, as nested dict
+ for cpu_layer_weight, gpu_layer_weight in zip(outputs_cpu['weights'], outputs_gpu['weights']):
+ for (cpu_weight, gpu_weight) in zip(cpu_layer_weight, gpu_layer_weight):
+ self.assertEqual(cpu_weight.grad, gpu_weight.grad, prec=5e-5)
+
+
+ input_val = torch.randn(seq_length, batch, input_size)
+ hx_val = torch.randn(num_layers, batch, hidden_size)
+ # FIXME: add bidirectional
+ # FIXME: add dropout
+ for module in (nn.RNN, nn.LSTM, nn.GRU):
+ for bias in (True, False):
+ rnn = module(input_size, hidden_size, num_layers, bias=bias)
+ outputs_cpu = forward_backward(False, rnn, input_val, hx_val, rnn.all_weights)
+
+ rnn_gpu = module(input_size, hidden_size, num_layers, bias=bias)
+ outputs_gpu = forward_backward(True, rnn_gpu, input_val, hx_val, rnn.all_weights)
+
+ compare_cpu_gpu(outputs_cpu, outputs_gpu)
+
+ for nonlinearity in ('tanh', 'relu'):
+ rnn = nn.rnn.RNN(input_size, hidden_size, num_layers, bias=bias, nonlinearity=nonlinearity)
+ outputs_cpu = forward_backward(False, rnn, input_val, hx_val, rnn.all_weights)
+
+ rnn_gpu = nn.rnn.RNN(input_size, hidden_size, num_layers, bias=bias, nonlinearity=nonlinearity)
+ outputs_gpu = forward_backward(True, rnn_gpu, input_val, hx_val, rnn.all_weights)
+
+ compare_cpu_gpu(outputs_cpu, outputs_gpu)
 
 
 def add_test(test):

torch/autograd/function.py

@@ -1,9 +1,8 @@
+import torch
 import torch._C as _C
 from collections import OrderedDict
 from itertools import chain
 
-import torch # FIXME: is this ok? Needed for torch.is_tensor
-import collections
 
 class Function(_C._FunctionBase):
 
@@ -78,12 +77,9 @@ def _iter(obj):
 _iter_tensors = _iter_filter(torch.is_tensor)
 _iter_None_tensors = _iter_filter(lambda o: o is None or torch.is_tensor(o))
 _map_variable_tensor = _nested_map(lambda o: isinstance(o, torch.autograd.Variable), lambda o: o.data)
-_map_tensor_type = _nested_map(lambda o: torch.is_tensor(o), lambda o: o.type())
 
 def _map_tensor_fromiter(itr):
- return _nested_map(lambda o: torch.is_tensor(o), lambda o: itr.next())
-def _map_variable_fromiter(itr):
- return _nested_map(lambda o: isinstance(o, torch.autograd.Variable), lambda o: itr.next())
+ return _nested_map(lambda o: torch.is_tensor(o), lambda o: next(itr))
 
 class NestedIOFunction(Function):
 
@@ -99,6 +95,7 @@ def backward(self, *gradients):
  nested_gradients = _map_tensor_fromiter(iter(gradients))(self._nested_output)
  del self._nested_output
  result = self.backward_extended(*nested_gradients)
+ del self._to_save_nested
  return tuple(_iter_None_tensors(result))
 
  __call__ = _do_forward
@@ -130,4 +127,3 @@ def forward_extended(self, *input):
 
  def backward_extended(self, *grad_output):
  raise NotImplementedError
- raise NotImplementedError

torch/autograd/functions/tensor.py

@@ -44,9 +44,15 @@ def forward(self, i, value=None):
  return i
 
  def backward(self, grad_output):
- grad_input = grad_output.clone()
- grad_input[self.index].fill_(0)
- return grad_input
+ if self.value is None:
+ grad_input = grad_output.clone()
+ grad_value = grad_output[self.index].clone()
+ grad_input[self.index].fill_(0)
+ return grad_input, grad_value
+ else:
+ grad_input = grad_output.clone()
+ grad_input[self.index].fill_(0)
+ return grad_input
 
 
 class Transpose(Function):

torch/autograd/variable.py

@@ -58,7 +58,8 @@ def __setitem__(self, key, value):
  return MaskedFill(value, inplace=True)(self, key)
  if isinstance(value, Variable):
  return SetItem(key)(self, value)
- return SetItem(key, value)(self)
+ else:
+ return SetItem(key, value)(self)
 
  def __iter__(self):
  return iter(map(lambda i: self[i], range(self.size(0))))

torch/backends/cudnn/__init__.py

@@ -280,7 +280,7 @@ def int_array(itr):
  return array_type(*itr)
 
 def descriptor(tensor, N=None):
- if N:
+ if N is not None:
  descriptor = TensorDescriptorArray(N)
  else:
  descriptor = TensorDescriptor()