torch.distributions.Categorical.sample uses unnecessary huge amount of memory

[jjabo]

🐛 Bug

When sampling from a torch.distributions.Categorical distribution initialized with a few distributions (e.g. 2 distributions each with 2**16 categories) the sample() functions allocates unnecessary huge amount of memory.

To Reproduce

import torch 
p = torch.ones([2, 2**16], device='cuda:0')
dist = torch.distributions.Categorical(p)
dist.sample([2**16])

Output:
Traceback (most recent call last):
File ".../lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3326, in run_code
exec(code_obj, self.user_global_ns, self.user_ns)
File "", line 4, in
dist.sample([2**16])
File ".../lib/python3.7/site-packages/torch/distributions/categorical.py", line 106, in sample
probs_2d = probs.reshape(-1, self._num_events)
RuntimeError: CUDA out of memory. Tried to allocate 32.00 GiB (GPU 0; 1.95 GiB total capacity; 1024.00 KiB already allocated; 1.15 GiB free; 2.00 MiB reserved in total by PyTorch)

Expected behavior

Looking into the torch.distributions.Categorical source shows:

    def sample(self, sample_shape=torch.Size()):
        sample_shape = self._extended_shape(sample_shape)
        param_shape = sample_shape + torch.Size((self._num_events,))
        probs = self.probs.expand(param_shape)
        probs_2d = probs.reshape(-1, self._num_events)
        sample_2d = torch.multinomial(probs_2d, 1, True)
        return sample_2d.reshape(sample_shape)

The last three lines allocate the new memory and use it. probs_2d becomes a potentially huge tensor with repeated information.
The following code snippet solves it and uses much less memory, but I'm not sure it covers all cases and whether the order of the dimensions is as expected. It's not thoroughly tested.

    def sample(self, sample_shape=torch.Size()):
        if not isinstance(sample_shape, torch.Size):
            sample_shape = torch.Size(sample_shape)
        samples = torch.multinomial(self.probs, sample_shape.numel(), True)
        return samples.reshape([-1] + list(sample_shape))

cc @vincentqb @fritzo @neerajprad @alicanb @vishwakftw @VitalyFedyunin @ngimel

[agolynski]

Thank you for the through diagnostic and code sample!

[neerajprad]

@jjabo - I think one issue is that the probs argument to Categorical can be batched (n-dimensional where n > 2, the last dim is taken as the event dim), whereas the argument to multinomial has to be 1 or 2D. I think the issue with the current code is the expand step which in this case expands probs to have shape more than 2**16 x 2**16! I think we should be able to use your solution to work for the batched case too as follows:

    def sample(self, sample_shape=torch.Size()):
        probs_2d = self.probs.reshape(-1, self._num_events)
        samples_2d = torch.multinomial(probs_2d, sample_shape.numel(), True).T
        return samples_2d.reshape(self._extended_shape(sample_shape))

Does that work on your test case?

[jjabo]

@neerajprad - The code you posted crashes for me with
AttributeError: 'list' object has no attribute 'numel'
since I give a list to sample_shape (is that wrong usage)? If I change your code to:

    def sample(self, sample_shape=torch.Size()):
        if not isinstance(sample_shape, torch.Size):
             sample_shape = torch.Size(sample_shape)
        probs_2d = self.probs.reshape(-1, self._num_events)
        samples_2d = torch.multinomial(probs_2d, sample_shape.numel(), True).T
        return samples_2d.reshape(self._extended_shape(sample_shape))

it works.
One peculiarity though is that the batch dimensions become the last dimensions of the output, e.g if I create categorical with probs.shape = [B,X] and call sample(sample_shape=[500]) the output has shape [500, B]. Is that as it should be? Shouldn't it be [B, 500]?

[fritzo]

@jjabo [500,B] is correct in your example since sample_shape is by convention left of batch_shape. From a tutorial:

      |      iid     | independent | dependent
------+--------------+-------------+------------
shape = sample_shape + batch_shape + event_shape

[neerajprad]

@jjabo - I have put up a fix with your solution above. I have not been able to test it on CUDA, but going by your response above, I think it should work and not result in OOM error anymore. Thanks for flagging this!