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
@@ -142,7 +142,7 @@ def sample(self):
     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))
+        return 0.0  # Works with cpu and cuda tensors.
 
 
 @contextlib.contextmanager
@@ -251,6 +251,8 @@ def map_data(name, data, fn, batch_size=0, batch_dim=0):
     if isinstance(data, (torch.Tensor, Variable)):
         size = data.size(batch_dim)
         with iarange(name, size, batch_size) as batch:
+            if data.is_cuda:
+                batch = batch.cuda()
             return fn(batch, data.index_select(batch_dim, batch))
     else:
         size = len(data)


Test patch: