SelfContactLoss - memory issues (running on GPU)

[GuyTevet]

Hi Lea:)
Applying SelfContactLoss with default parameters on batch_size=32 (i.e. [32, 10475, 3] tensor for a smplx mesh), fails for trying to allocate 112GB of GPU memory. Is there a way to handle this?

RuntimeError: CUDA out of memory. Tried to allocate 112.33 GiB (GPU 0; 10.76 GiB total capacity; 309.89 MiB already allocated; 9.30 GiB free; 346.00 MiB reserved in total by PyTorch)

[GuyTevet]

The memory issue occurs in the winding_numbers module:

Traceback (most recent call last):
  File "/disk2/guytevet/.pycharm_helpers/pydev/_pydevd_bundle/pydevd_exec2.py", line 3, in Exec
    exec(exp, global_vars, local_vars)
  File "<input>", line 1, in <module>
  File "/disk2/guytevet/venvs/py3.7_vibe/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/disk2/guytevet/selfcontact/selfcontact/losses/general_selfcontact_loss.py", line 80, in forward
    test_segments=self.test_segments
  File "/disk2/guytevet/selfcontact/selfcontact/selfcontact.py", line 198, in segment_vertices
    test_segments
  File "/disk2/guytevet/selfcontact/selfcontact/selfcontact.py", line 136, in get_intersection_mask
    triangles).le(0.99)
  File "/disk2/guytevet/selfcontact/selfcontact/utils/mesh.py", line 224, in winding_numbers
    points, triangles, thresh=thresh).sum(dim=-1)
  File "/disk2/guytevet/selfcontact/selfcontact/utils/mesh.py", line 155, in solid_angles
    centered_tris = triangles[:, None] - points[:, :, None, None]
RuntimeError: CUDA out of memory. Tried to allocate 112.33 GiB (GPU 0; 10.76 GiB total capacity; 309.89 MiB already allocated; 9.30 GiB free; 346.00 MiB reserved in total by PyTorch)

[muelea]

Hi Guy,

the code was tested with 16 GB of RAM.

Since your RAM size is 12 GB, you can try to split the number of points to be tested into n chunks and do the intersection tests n times. This way the intersection tests will need less memory, but more time. The code block you want to change is here: https://github.com/muelea/selfcontact/blob/master/selfcontact/selfcontact.py#L135

For example, to replace two by four blocks, replace

# split because of memory into two chunks
exterior = torch.zeros((bs, nv), device=vertices.device,
    dtype=torch.bool)
exterior[:, :5000] = winding_numbers(vertices[:,:5000,:],
    triangles).le(0.99)
exterior[:, 5000:] = winding_numbers(vertices[:,5000:,:],
    triangles).le(0.99)

by

# split because of memory into four chunks
exterior = torch.zeros((bs, nv), device=vertices.device,
    dtype=torch.bool)
exterior[:, :3000] = winding_numbers(vertices[:,:3000,:],
    triangles).le(0.99)
exterior[:, 3000:6000] = winding_numbers(vertices[:,3000:6000,:],
    triangles).le(0.99)
exterior[:, 6000:9000] = winding_numbers(vertices[:,6000:9000,:],
    triangles).le(0.99)
exterior[:, 9000:] = winding_numbers(vertices[:,9000:,:],
    triangles).le(0.99)

If you use HD vertices, make the same changes here as well: https://github.com/muelea/selfcontact/blob/master/selfcontact/selfcontact.py#L161

Hope that helps,
Lea

[GuyTevet]

Cool! Thanks:)