优化高斯热力图生成：消除循环及1D高斯实现问题排查

问题描述 #

我有一组尺寸为[68, 2]的关键点，需将其映射为高斯热力图，为此编写了如下函数：

def generate_gaussian(t, x, y, sigma=10):
    """
    Generates a 2D Gaussian point at location x,y in tensor t.
    x should be in range (-1, 1).
    sigma is the standard deviation of the generated 2D Gaussian.
    """
    h,w = t.shape

    # Heatmap pixel per output pixel
    mu_x = int(0.5 * (x + 1.) * w)
    mu_y = int(0.5 * (y + 1.) * h)

    tmp_size = sigma * 3

    # Top-left
    x1,y1 = int(mu_x - tmp_size), int(mu_y - tmp_size)

    # Bottom right
    x2, y2 = int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)
    if x1 >= w or y1 >= h or x2 < 0 or y2 < 0:
        return t

    size = 2 * tmp_size + 1
    tx = np.arange(0, size, 1, np.float32)
    ty = tx[:, np.newaxis]
    x0 = y0 = size // 2

    # The gaussian is not normalized, we want the center value to equal 1
    g = torch.tensor(np.exp(- ((tx - x0) ** 2 + (ty - y0) ** 2) / (2 * sigma ** 2)))

    # Determine the bounds of the source gaussian
    g_x_min, g_x_max = max(0, -x1), min(x2, w) - x1
    g_y_min, g_y_max = max(0, -y1), min(y2, h) - y1

    # Image range
    img_x_min, img_x_max = max(0, x1), min(x2, w)
    img_y_min, img_y_max = max(0, y1), min(y2, h)

    t[img_y_min:img_y_max, img_x_min:img_x_max] = \
      g[g_y_min:g_y_max, g_x_min:g_x_max]

    return t

def rescale(a, img_size):
    # scale tensor to [-1, 1]
    return 2 * a / img_size[0] - 1

目前我通过循环为每个关键点计算高斯热力图，再堆叠为[68, H, W]张量：

x_k1 = [generate_gaussian(torch.zeros(H, W), x, y) for x, y in rescale(kp1.numpy(), frame.shape)]
x_k1 = torch.stack(x_k1, dim=0)

但该方法速度极慢，想知道有没有办法去除循环？

此外，我尝试了基于1D高斯的实现方案：

def generate_gaussian1D(t, x, y, sigma=10):
    h,w = t.shape

    # Heatmap pixel per output pixel
    mu_x = int(0.5 * (x + 1.) * w)
    mu_y = int(0.5 * (y + 1.) * h)

    tmp_size = sigma * 3

    # Top-left
    x1, y1 = int(mu_x - tmp_size), int(mu_y - tmp_size)

    # Bottom right
    x2, y2 = int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)
    if x1 >= w or y1 >= h or x2 < 0 or y2 < 0:
        return t

    size = 2 * tmp_size + 1
    tx = np.arange(0, size, 1, np.float32)
    ty = tx[:, np.newaxis]
    x0 = y0 = size // 2

    g = torch.tensor(np.exp(-np.power(tx - mu_x, 2.) / (2 * np.power(sigma, 2.))))
    g = g * g[:, None]

    g_x_min, g_x_max = max(0, -x1), min(x2, w) - x1
    g_y_min, g_y_max = max(0, -y1), min(y2, h) - y1

    img_x_min, img_x_max = max(0, x1), min(x2, w)
    img_y_min, img_y_max = max(0, y1), min(y2, h)

    t[img_y_min:img_y_max, img_x_min:img_x_max] = \
      g[g_y_min:g_y_max, g_x_min:g_x_max]

    return t

但输出的高斯热力图不完整（如图所示），不确定问题出在哪里，恳请帮助。

解决方案 #

一、批量生成热力图，彻底移除循环 #

循环慢的核心是没用到PyTorch的向量化运算优势，直接批量生成[68, H, W]张量即可，步骤如下：

1. 批量转换关键点坐标
   一次性把所有[68,2]的关键点从[-1,1]映射到热力图像素坐标：

   kp_scaled = rescale(kp1, frame.shape)  # shape [68, 2]
   W, H = frame.shape[1], frame.shape[0]  # 注意frame的维度顺序
   mu_x = (0.5 * (kp_scaled[:, 0] + 1.) * W).long()
   mu_y = (0.5 * (kp_scaled[:, 1] + 1.) * H).long()
   

2. 生成热力图网格
   创建x/y方向的网格坐标，通过维度扩展适配批量广播：

   y_grid, x_grid = torch.meshgrid(torch.arange(H), torch.arange(W), indexing='ij')
   # 扩展维度：[H,W] → [1,H,W]，和[68,1,1]的mu做广播运算
   x_grid = x_grid.unsqueeze(0)
   y_grid = y_grid.unsqueeze(0)
   

3. 批量计算高斯热力图
   利用广播机制一次性生成所有关键点的热力图：

   sigma = 10
   # 直接计算所有68个热力图的高斯值
   gaussian = torch.exp(
       -((x_grid - mu_x.unsqueeze(-1).unsqueeze(-1))**2 + 
         (y_grid - mu_y.unsqueeze(-1).unsqueeze(-1))**2) 
       / (2 * sigma**2)
   )
   

   得到的gaussian就是[68, H, W]的目标张量，速度比循环快一个数量级以上。

   如果要保留原代码中“超出边界不生成高斯”的逻辑，可添加掩码处理：

   tmp_size = sigma * 3
   # 计算每个关键点的边界范围
   x1 = mu_x - tmp_size
   x2 = mu_x + tmp_size + 1
   y1 = mu_y - tmp_size
   y2 = mu_y + tmp_size + 1
   # 标记有效关键点（未超出热力图边界）
   valid_mask = (x1 < W) & (y1 < H) & (x2 > 0) & (y2 > 0)
   # 无效关键点对应的热力图置0
   gaussian[~valid_mask] = 0.
   

二、修复1D高斯实现的错误 #

你的generate_gaussian1D函数问题出在高斯核的计算逻辑：
原代码中用tx - mu_x计算偏移，但tx是高斯核内部的相对坐标（0到size-1），mu_x是热力图的绝对像素坐标，两者相减完全错误。正确做法是基于核中心x0计算相对偏移：

修改后的高斯核计算部分：

# 替换原错误的g计算代码
g_x = torch.tensor(np.exp(-np.power(tx - x0, 2.) / (2 * np.power(sigma, 2.))))
g_y = torch.tensor(np.exp(-np.power(tx - x0, 2.) / (2 * np.power(sigma, 2.))))
g = g_y.unsqueeze(1) * g_x  # 外积生成2D高斯核

这样生成的高斯核中心值为1，分布完全符合预期，不会出现截断或错位的问题。

内容的提问来源于stack exchange，提问作者monopoly