Why index_points() use so much memory?

username: haoran-hash
cooked: <p>code snippet is below.</p>
<pre><code class="lang-auto">import torch
from torch.profiler import profile, ProfilerActivity
import logging

def index_points(points, indices):
    """
    points.shape = (b, n, c)
    indices.shape = (b, m, k)
    return res.shape = (b, m, k, c)
    """
    _, _, c = points.shape
    _, m, _ = indices.shape

    indices = indices.unsqueeze(dim=3).expand(-1, -1, -1, c)
    points = points.unsqueeze(dim=1).expand(-1, m, -1, -1)
    res = points.gather(dim=2, index=indices)

    return res

def index_points_2(points, indices):
    """
    points.shape = (b, n, c)
    indices.shape = (b, m, k)
    return res.shape = (b, m, k, c)
    """
    device = points.device
    b, m, k = indices.shape

    batch_indices = torch.arange(b, device=device).view(b, 1, 1).expand(-1, m, k)
    res = points[batch_indices, indices, :]
    return res

if __name__ == '__main__':
    b, n, m, k, c = 16, 512, 128, 64, 128
    features = torch.randn((b, n, c), device='cuda:0', requires_grad=True)
    indices = torch.randint(0, n, (b, m, k), device='cuda:0')
    
    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], profile_memory=True) as prof:
        res = index_points_2(features, indices)
        res.backward(torch.ones_like(res))
    
    logging.basicConfig(filename='test.log', format='%(message)s', level=logging.INFO)
    logging.info(prof.key_averages().table(sort_by='self_cuda_memory_usage', row_limit=20))
</code></pre>
<p><code>index_points()</code> and <code>index_points_2()</code> do the same thing, so they’re interchangeable.<br>
The results are as follows. Only the relevant columns are showed here<br>
<code>index_points()</code> result:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/8/88be87a83338939bc2108bcc8b72de7d7d9a55c7.png" data-download-href="https://discuss.pytorch.org/uploads/default/88be87a83338939bc2108bcc8b72de7d7d9a55c7" title="屏幕截图 2022-03-01 193610"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/8/88be87a83338939bc2108bcc8b72de7d7d9a55c7_2_609x500.png" alt="屏幕截图 2022-03-01 193610" data-base62-sha1="jvH8Xc03Cke6i9pE3yZt1ESgibR" width="609" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/8/88be87a83338939bc2108bcc8b72de7d7d9a55c7_2_609x500.png, https://discuss.pytorch.org/uploads/default/original/3X/8/8/88be87a83338939bc2108bcc8b72de7d7d9a55c7.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/8/8/88be87a83338939bc2108bcc8b72de7d7d9a55c7.png 2x" data-dominant-color="F6F4F4"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">屏幕截图 2022-03-01 193610</span><span class="informations">817×670 295 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<code>index_points_2()</code> result:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/1/0126d41207adcd43ae570157efba50c8d1e72ee9.png" data-download-href="https://discuss.pytorch.org/uploads/default/0126d41207adcd43ae570157efba50c8d1e72ee9" title="屏幕截图 2022-03-01 195140"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/1/0126d41207adcd43ae570157efba50c8d1e72ee9_2_610x500.png" alt="屏幕截图 2022-03-01 195140" data-base62-sha1="abFlH3lcEHq5IZCGNqOjE9N7kJ" width="610" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/1/0126d41207adcd43ae570157efba50c8d1e72ee9_2_610x500.png, https://discuss.pytorch.org/uploads/default/original/3X/0/1/0126d41207adcd43ae570157efba50c8d1e72ee9.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/0/1/0126d41207adcd43ae570157efba50c8d1e72ee9.png 2x" data-dominant-color="F7F5F5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">屏幕截图 2022-03-01 195140</span><span class="informations">825×676 285 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
From above image, we can know <code>index_points()</code> use more memory. But I do not konw why? Can anyone answer?</p>