Question about extreme memory fragmentation in PointNet++ sampling

eliphatfs: <p>I am meeting <em>EXTREME</em> fragmentation in my training, and I think it is due to the PointNet++ sample and group process here, because it was originally followed by <code>torch.cuda.empty_cache()</code> calls, but this dramatically slows down training and thus I removed them. I wonder why this happens – number of points and batch size are always kept the same in my training process, and the allocator should work well.</p>
<p>The code is modified from <a href="https://github.com/yanx27/Pointnet_Pointnet2_pytorch/tree/master/models" rel="noopener nofollow ugc">Pointnet_Pointnet2_pytorch/models at master · yanx27/Pointnet_Pointnet2_pytorch · GitHub</a></p>
<pre><code class="lang-python">def square_distance(src, dst):
    """
    Calculate Euclid distance between each two points.

    src^T * dst = xn * xm + yn * ym + zn * zm；
    sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
    sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
    dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
         = sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst

    Input:
        src: source points, [B, N, C]
        dst: target points, [B, M, C]
    Output:
        dist: per-point square distance, [B, N, M]
    """
    B, N, _ = src.shape
    _, M, _ = dst.shape
    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
    dist += torch.sum(src ** 2, -1).view(B, N, 1)
    dist += torch.sum(dst ** 2, -1).view(B, 1, M)
    return dist


def index_points(points, idx):
    """

    Input:
        points: input points data, [B, N, C]
        idx: sample index data, [B, S]
    Return:
        new_points:, indexed points data, [B, S, C]
    """
    device = points.device
    B = points.shape[0]
    view_shape = list(idx.shape)
    view_shape[1:] = [1] * (len(view_shape) - 1)
    repeat_shape = list(idx.shape)
    repeat_shape[0] = 1
    batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape)
    new_points = points[batch_indices, idx, :]
    return new_points


def farthest_point_sample(xyz, npoint):
    """
    Input:
        xyz: pointcloud data, [B, N, 3]
        npoint: number of samples
    Return:
        centroids: sampled pointcloud index, [B, npoint]
    """
    return dgl.geometry.farthest_point_sampler(xyz, npoint)


def query_ball_point(radius, nsample, xyz, new_xyz):
    """
    Input:
        radius: local region radius
        nsample: max sample number in local region
        xyz: all points, [B, N, 3]
        new_xyz: query points, [B, S, 3]
    Return:
        group_idx: grouped points index, [B, S, nsample]
    """
    device = xyz.device
    B, N, C = xyz.shape
    _, S, _ = new_xyz.shape
    group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N).repeat([B, S, 1])
    sqrdists = square_distance(new_xyz, xyz)
    group_idx[sqrdists &gt; radius ** 2] = N
    group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample]
    group_first = group_idx[..., :1].repeat([1, 1, nsample])
    mask = group_idx == N
    group_idx[mask] = group_first[mask]
    return group_idx


def sample_and_group(npoint, radius, nsample, xyz, points, returnfps=False):
    """
    Input:
        npoint:
        radius:
        nsample:
        xyz: input points position data, [B, N, 3]
        points: input points data, [B, N, D]
    Return:
        new_xyz: sampled points position data, [B, npoint, nsample, 3]
        new_points: sampled points data, [B, npoint, nsample, 3+D]
    """
    B, N, C = xyz.shape
    S = npoint
    fps_idx = farthest_point_sample(xyz, npoint) # [B, npoint, C]
    # torch.cuda.empty_cache()
    new_xyz = index_points(xyz, fps_idx)
    # torch.cuda.empty_cache()
    idx = query_ball_point(radius, nsample, xyz, new_xyz)
    # torch.cuda.empty_cache()
    grouped_xyz = index_points(xyz, idx) # [B, npoint, nsample, C]
    # torch.cuda.empty_cache()
    grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C)
    # torch.cuda.empty_cache()

    if points is not None:
        grouped_points = index_points(points, idx)
        new_points = torch.cat([grouped_xyz_norm, grouped_points], dim=-1) # [B, npoint, nsample, C+D]
    else:
        new_points = grouped_xyz_norm
    if returnfps:
        return new_xyz, new_points, grouped_xyz, fps_idx
    else:
        return new_xyz, new_points
</code></pre>
<p>By <em>EXTREME</em> fragmentation I mean:</p>
<pre><code class="lang-plaintext">  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File ".../src/models/pointnet_util.py", line 199, in forward
    new_xyz, new_points = sample_and_group(self.npoint, self.radius, self.nsample, xyz, points)
  File ".../src/models/pointnet_util.py", line 130, in sample_and_group
    idx = query_ball_point(radius, nsample, xyz, new_xyz)
  File ".../src/models/pointnet_util.py", line 105, in query_ball_point
    group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample]
RuntimeError: CUDA out of memory. Tried to allocate 7.63 GiB (GPU 0; 79.35 GiB total capacity; 27.47 GiB already allocated; 5.05 GiB free; 72.49 GiB reserved in total by PyTorch) If reserved memory is &gt;&gt; allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
</code></pre>
<p>It is like crazy – 27GB allocated but 72GB reserved.</p>
<p>Any input is welcome, and thanks in advance!</p>