Repository: pyro-ppl/pyro

Issue:
CUDA Support for map_data
Using `map_data` when the `data` tensor is on the GPU doesn't work at the moment. It looks like the problem stems from the fact that `sample` and `batch_log_pdf` of the `Subsample` distribution always return cpu tensors. The former is a problem when using the indices to index into gpu data. The latter when accumulating `log_pdf` in the trace when other choices return gpu values for their `log_pdf`.

AFAICT we can't infer whether these things should happen on the cpu/gpu from `Subsample`'s parameters, so perhaps we need to it pass an extra `use_cuda` flag? This could be set automatically from `map_data` (by looking at the `data` tensor), but would need to be passed explicitly when calling `iarange` directly, which is a drag. Hopefully there's something better we can do.

One vague idea is that since accumulating `log_pdf` on the trace requires all choices to return log probs of the same type, perhaps it's reasonable to say at the level of a trace whether log probs should be computed on the gpu or cpu, and to pass this info on to distributions so they can use it to determine appropriate return types, rather than leaving each distribution to figure this out independently as at present. (This doesn't seem right, but I'll leave it here in case it's useful.)


Discussion:
Good catch!  We definitely need a dedicated CUDA CI machine to catch things like this.  

The easiest hack to get around this issue at the moment would be to check PyTorch's global default tensor type and type the indices based on that.
Do you have a sense of how often we'd perform batched inference of a dataset that lives on GPU? In my probably outdated experience, GPU memory pressure is the main reason to batch data, and so I'd think we would want to use custom batching for GPU-based computation anyway, maybe like the `iarange(size, subsample=custom_indices_that_live_on_gpu)` as in #258 .
> Do you have a sense of how often we'd perform batched inference of a dataset that lives on GPU?

@fritzo I'm not sure you were asking this of me, but fwiw, I don't personally have a good sense of how often this comes ups.

> and so I'd think we would want to use custom batching for GPU-based computation anyway

This sounds to me like `map_data` would be cpu only, which would be a shame, but I'm guessing I misunderstand. (ETA: And there'd still be the `log_pdf` part of the problem?)
> The easiest hack to get around this issue at the moment would be to check PyTorch's global default tensor type and type the indices based on that.

@eb8680 Thanks. When I set the default to some cuda type, I break `torch.ones` and friends.(`ng_ones`, etc.)

```
>>> torch.set_default_tensor_type('torch.cuda.FloatTensor')
>>> torch.ones(1)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: Type torch.cuda.FloatTensor doesn't implement stateless methods
```

Am I missing something?
@null-a Eli's approach sounds good, but I was suggesting another CUDA workaround like:
```py
def model(batch):
    with iarange(full_data_size, subsample=batch):
        # ...do stuff with batch...

inference = SVI(model, ...)
for cpu_batch in my_custom_batcher(batch_size):
    gpu_batch = cpu_batch.cuda()
    inference.step(data)
```
@rohitsingh0812 had this issue with `torch.eye`. There are a few tensor creation methods (like range, eye) which did not have cuda support (see https://github.com/pytorch/pytorch/issues/2165) until recently. The PRs for supporting these are merged in master, but I am not sure if they are on the latest release. Once it is, we should ideally just be able to specify the default tensor type as cuda, and everything should work as expected.
OK, as I understand, the root of this problem is that `_Subsample.log_pdf()` returns a `Variable(torch.zeros(1))`. The more minor issue is that `batch` is a cpu tensor, and this can be solved with a single line `batch = batch.cuda()`.

Also, apologies for introducing this breakage in #239 . Hopefully this kind of breakage will be addressed by testing cuda support in examples (#304) and by adding `--cuda` flags to the `examples/*.py` (and as we *commit early and often* :wink: ).

Resolving patch:
diff --git a/pyro/__init__.py b/pyro/__init__.py
--- a/pyro/__init__.py
+++ b/pyro/__init__.py
@@ -131,13 +131,14 @@ class _Subsample(Distribution):
     Internal use only. This should only be used by `iarange`.
     """
 
-    def __init__(self, size, subsample_size):
+    def __init__(self, size, subsample_size, use_cuda=False):
         """
         :param int size: the size of the range to subsample from
         :param int subsample_size: the size of the returned subsample
         """
         self.size = size
         self.subsample_size = subsample_size
+        self.use_cuda = use_cuda
 
     def sample(self):
         """
@@ -145,16 +146,18 @@ def sample(self):
         :rtype: torch.autograd.Variable of torch.LongTensor
         """
         assert 0 <= self.subsample_size <= self.size
-        return Variable(torch.randperm(self.size)[:self.subsample_size])
+        result = Variable(torch.randperm(self.size)[:self.subsample_size])
+        return result.cuda() if self.use_cuda else result
 
     def batch_log_pdf(self, x):
         # This is zero so that iarange can provide an unbiased estimate of
         # the non-subsampled batch_log_pdf.
-        return Variable(torch.zeros(0))
+        result = Variable(torch.zeros(1))
+        return result.cuda() if self.use_cuda else result
 
 
 @contextlib.contextmanager
-def iarange(name, size, subsample_size=0, subsample=None):
+def iarange(name, size, subsample_size=0, subsample=None, use_cuda=False):
     """
     Context manager for ranges indexing iid variables, optionally subsampling.
 
@@ -171,6 +174,7 @@ def iarange(name, size, subsample_size=0, subsample=None):
     :param subsample: Optional custom subsample for user-defined subsampling schemes.
         If specified, then `subsample_size` will be set to `len(subsample)`.
     :type subsample: Anything supporting `len()`.
+    :param bool use_cuda: Whether to use cuda tensors for `subsample` and `log_pdf`.
     :return: A context manager yielding a single 1-dimensional `torch.Tensor` of indices.
 
     Examples::
@@ -192,16 +196,18 @@ def iarange(name, size, subsample_size=0, subsample=None):
     """
     if subsample is not None:
         subsample_size = len(subsample)
+        use_cuda = use_cuda or getattr(subsample, 'is_cuda', False)
         assert subsample_size <= size, 'subsample is larger than size'
     elif subsample_size == 0 or subsample_size >= size:
         subsample_size = size
     if subsample_size == size:
         # If not subsampling, there is no need to scale and we can ignore the _PYRO_STACK.
-        yield Variable(torch.LongTensor(list(range(size))))
+        result = Variable(torch.LongTensor(list(range(size))))
+        yield result.cuda() if use_cuda else result
         return
 
     if subsample is None:
-        subsample = sample(name, _Subsample(size, subsample_size))
+        subsample = sample(name, _Subsample(size, subsample_size, use_cuda))
     if len(_PYRO_STACK) == 0:
         yield subsample
     else:
@@ -211,7 +217,7 @@ def iarange(name, size, subsample_size=0, subsample=None):
             yield subsample
 
 
-def irange(name, size, subsample_size=0, subsample=None):
+def irange(name, size, subsample_size=0, subsample=None, use_cuda=False):
     """
     Non-vectorized version of `iarange`. See `iarange` for details.
 
@@ -221,6 +227,7 @@ def irange(name, size, subsample_size=0, subsample=None):
     :param subsample: Optional custom subsample for user-defined subsampling schemes.
         If specified, then `subsample_size` will be set to `len(subsample)`.
     :type subsample: Anything supporting `len()`.
+    :param bool use_cuda: Whether to use cuda tensors for `subsample` and `log_pdf`.
     :return: A context manager yielding a single 1-dimensional `torch.Tensor` of indices.
 
     Examples::
@@ -232,7 +239,7 @@ def irange(name, size, subsample_size=0, subsample=None):
     if subsample is not None:
         subsample_size = len(subsample)
 
-    with iarange(name, size, subsample_size, subsample) as batch:
+    with iarange(name, size, subsample_size, subsample, use_cuda) as batch:
         # Wrap computation in an independence context.
         indep_context = LambdaPoutine(None, name, 1.0, 'list', 0, subsample_size)
         if isinstance(batch, Variable):
@@ -242,7 +249,7 @@ def irange(name, size, subsample_size=0, subsample=None):
                 yield i
 
 
-def map_data(name, data, fn, batch_size=0, batch_dim=0):
+def map_data(name, data, fn, batch_size=0, batch_dim=0, use_cuda=False):
     """
     Data subsampling with the important property that all the data are conditionally independent.
 
@@ -254,15 +261,17 @@ def map_data(name, data, fn, batch_size=0, batch_dim=0):
     :param callable fn: a function taking `(index, datum)` pairs, where `dataum = data[index]`
     :param int batch_size: number of samples per batch, or zero for the entire dataset
     :param int batch_dim: dimension to subsample for tensor inputs
+    :param bool use_cuda: Whether to use cuda tensors for `subsample` and `log_pdf`.
     :return: a list of values returned by `fn`
     """
+    use_cuda = use_cuda or getattr(data, 'is_cuda', False)
     if isinstance(data, (torch.Tensor, Variable)):
         size = data.size(batch_dim)
-        with iarange(name, size, batch_size) as batch:
+        with iarange(name, size, batch_size, use_cuda=use_cuda) as batch:
             return fn(batch, data.index_select(batch_dim, batch))
     else:
         size = len(data)
-        return [fn(i, data[i]) for i in irange(name, size, batch_size)]
+        return [fn(i, data[i]) for i in irange(name, size, batch_size, use_cuda=use_cuda)]
 
 
 # XXX this should have the same call signature as torch.Tensor constructors
diff --git a/pyro/distributions/uniform.py b/pyro/distributions/uniform.py
--- a/pyro/distributions/uniform.py
+++ b/pyro/distributions/uniform.py
@@ -57,10 +57,10 @@ def batch_log_pdf(self, x, a=None, b=None, *args, **kwargs):
         if x.size != a.size():
             a = a.expand_as(x)
             b = b.expand_as(x)
-        l = x.ge(a).type_as(a)
-        u = x.le(b).type_as(b)
+        lb = x.ge(a).type_as(a)
+        ub = x.le(b).type_as(b)
         batch_log_pdf_shape = self.batch_shape(a, b) + (1,)
-        return torch.sum(torch.log(l.mul(u)) - torch.log(b - a), -1).contiguous().view(batch_log_pdf_shape)
+        return torch.sum(torch.log(lb.mul(ub)) - torch.log(b - a), -1).contiguous().view(batch_log_pdf_shape)
 
     def analytic_mean(self, a=None, b=None):
         a, b = self._sanitize_input(a, b)


Test patch:
diff --git a/tests/poutine/test_mapdata.py b/tests/poutine/test_mapdata.py
--- a/tests/poutine/test_mapdata.py
+++ b/tests/poutine/test_mapdata.py
@@ -9,7 +9,7 @@
 import pyro.poutine as poutine
 import pyro.optim as optim
 from pyro.infer import SVI
-from tests.common import assert_equal
+from tests.common import assert_equal, requires_cuda
 
 
 @pytest.mark.stage("integration", "integration_batch_1")
@@ -216,3 +216,46 @@ def test_custom_subsample(model):
     subsample = [1, 3, 5, 7]
     assert model(subsample) == subsample
     assert poutine.trace(model)(subsample) == subsample
+
+
+def iarange_cuda_model():
+    mu = Variable(torch.zeros(20).cuda())
+    sigma = Variable(torch.ones(20).cuda())
+    with pyro.iarange("data", 20, 5, use_cuda=True) as batch:
+        pyro.sample("x", dist.diagnormal, mu[batch], sigma[batch])
+
+
+def irange_cuda_model():
+    mu = Variable(torch.zeros(20).cuda())
+    sigma = Variable(torch.ones(20).cuda())
+    for i in pyro.irange("data", 20, 5, use_cuda=True):
+        pyro.sample("x_{}".format(i), dist.diagnormal, mu[i], sigma[i])
+
+
+def map_data_vector_cuda_model():
+    mu = Variable(torch.zeros(20).cuda())
+    sigma = Variable(torch.ones(20).cuda())
+    pyro.map_data("data", mu,
+                  lambda i, mu: pyro.sample("x", dist.diagnormal, mu, sigma[i]),
+                  use_cuda=True)
+
+
+def map_data_iter_cuda_model():
+    mu = Variable(torch.zeros(20).cuda())
+    sigma = Variable(torch.ones(20).cuda())
+    pyro.map_data("data", list(mu),
+                  lambda i, mu: pyro.sample("x_{}".format(i), dist.diagnormal, mu, sigma[i]),
+                  use_cuda=True)
+
+
+@requires_cuda
+@pytest.mark.parametrize('model', [
+    iarange_cuda_model,
+    irange_cuda_model,
+    map_data_vector_cuda_model,
+    map_data_iter_cuda_model,
+], ids=["iarange", "irange", "map_data_vector", "map_data_iter"])
+def test_cuda(model):
+    tr = poutine.trace(model).get_trace()
+    assert tr.log_pdf().is_cuda
+    assert tr.batch_log_pdf().is_cuda