Question about extreme memory fragmentation in PointNet++ sampling

I am meeting EXTREME fragmentation in my training, and I think it is due to the PointNet++ sample and group process here, because it was originally followed by torch.cuda.empty_cache() 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.

The code is modified from Pointnet_Pointnet2_pytorch/models at master · yanx27/Pointnet_Pointnet2_pytorch · GitHub

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

        src: source points, [B, N, C]
        dst: target points, [B, M, C]
        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):

        points: input points data, [B, N, C]
        idx: sample index data, [B, S]
        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):
        xyz: pointcloud data, [B, N, 3]
        npoint: number of samples
        centroids: sampled pointcloud index, [B, npoint]
    return dgl.geometry.farthest_point_sampler(xyz, npoint)

def query_ball_point(radius, nsample, xyz, new_xyz):
        radius: local region radius
        nsample: max sample number in local region
        xyz: all points, [B, N, 3]
        new_xyz: query points, [B, S, 3]
        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 > 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):
        xyz: input points position data, [B, N, 3]
        points: input points data, [B, N, D]
        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 =[grouped_xyz_norm, grouped_points], dim=-1) # [B, npoint, nsample, C+D]
        new_points = grouped_xyz_norm
    if returnfps:
        return new_xyz, new_points, grouped_xyz, fps_idx
        return new_xyz, new_points

By EXTREME fragmentation I mean:

  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File ".../src/models/", line 199, in forward
    new_xyz, new_points = sample_and_group(self.npoint, self.radius, self.nsample, xyz, points)
  File ".../src/models/", line 130, in sample_and_group
    idx = query_ball_point(radius, nsample, xyz, new_xyz)
  File ".../src/models/", 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 >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

It is like crazy – 27GB allocated but 72GB reserved.

Any input is welcome, and thanks in advance!