Open gshartnett opened 3 years ago
As a quick workaround, you could set cuda as default before creating your spline_coupling
object
torch.set_default_tensor_type(torch.cuda.FloatTensor)
spline_transform = T.spline_coupling(...)
Another not-ideal workaround is to save-and-load with map_location:
torch.save(any_python_object, "temp_file.pt")
any_python_object = torch.load("temp_file", map_location="cuda:0")
Hi @fritzo , I tried the first solution (torch.set_default_tensor_type(torch.cuda.FloatTensor)
) and it doesn't seem to work. I get this error when using torch.DataLoaders
and pytroch_lightning.LightningDataModule
RuntimeError: Expected a 'cuda' device type for generator but found 'cpu'
Is there some other workaround to try out? Also, what's the bottleneck here? Looking at the traceback it fails at log_prob()
evaluation.
Hi @gshartnett, it's difficult to diagnose where the stray cpu tensor is coming from without actually diving into a debugger. I'd recommend running under pdb and inspecting each tensor's device. You might also try updating PyTorch, since they seem to be getting more consistent over time.
One thing you could try is to use torch.save and torch.load, but you'd only be able to do that once before training:
x = my_complex_data_structure()
torch.save(x, "temp.pt")
x = torch.load("temp.pt", map_location="cuda:0")
thanks for the very prompt reply @fritzo , I was able to send some flows to cuda
with self.to(device)
.
E.g.
AffineAutoregressive(
AutoRegressiveNN(
latent_dim,
[hidden_units for _ in range(n_hidden)],
skip_connections=True,
)
).to(device)
Amongst the one I tried, this also works for BatchNorm
and AffineCoupling
but it does not work for Permute
(which returns an error like "permute does not have to() method"). From very high level inspection, Permute
is the only one inheriting from torch.distributions.transforms.Transform
instead of pyro.distributions.torch_transform.TransformModule
amongst the one I tried (the latter inheriting from nn.Module
). Might be off track, but anyway wanted to report this.
EDIT:
I was wrong, subclassing with TransformModule
doesn't solve the problem. It seems to be in
return y.index_select(self.dim, self.inv_permutation)
SOLUTION: simply send to device the permutation index
Permute(torch.LongTensor(perm).to(device))
thanks again for the help!
I am interested in using the Pyro implementation of normalizing flows for my research. However, I cannot find anywhere in the docs instructions on how to enable GPU support. The example page makes no mention of GPUs, and if I modify that code from
to
and attempt to run, I get this error
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
. How should I move the TransformedDistributions class entirely onto the GPU device? Tagging @stefanwebb since he is the person responsible for normalizing flow support according to this issue.