如何在PyTorch中为给定点张量在特征图上生成二维高斯分布

PyTorch没有直接封装该需求的官方API，但可以通过基础张量操作快速实现，性能可以满足大部分场景需求。

核心实现思路 #

• 先确定高斯核参数：高斯作用半径r、标准差sigma、输出特征图的尺寸H x W
• 生成尺寸为(2r+1)x(2r+1)的归一化二维高斯核模板
• 遍历所有输入坐标点，将高斯核模板贴到特征图对应位置，重叠区域直接累加
• 坐标数量很大时可以用卷积操作替代循环，进一步提升效率

代码实现 #

1. 生成二维高斯核 #

import torch
import torch.nn.functional as F

def generate_2d_gaussian_kernel(kernel_size, sigma=None):
    # kernel_size需为奇数，如3、5、7
    k = kernel_size // 2
    x, y = torch.meshgrid(torch.arange(-k, k+1), torch.arange(-k, k+1), indexing='ij')
    if sigma is None:
        # 自动计算sigma的通用经验值
        sigma = 0.3 * (k - 1) + 0.8
    gaussian = torch.exp(-(x**2 + y**2) / (2 * sigma**2))
    # 可选归一化规则：除以总和保证核总和为1，或除以最大值保证峰值为1
    gaussian = gaussian / gaussian.sum()
    return gaussian

2. 少量坐标点场景：直接贴核实现 #

适合坐标点数量<1000的场景，避免卷积的额外开销：

def draw_gaussian_on_featuremap(points, H, W, kernel_radius=3, sigma=None):
    # points形状为(N, 2)，每个元素是[x, y]坐标，x对应W维度，y对应H维度
    feat_map = torch.zeros((H, W), device=points.device)
    kernel_size = 2 * kernel_radius + 1
    gaussian_kernel = generate_2d_gaussian_kernel(kernel_size, sigma).to(points.device)
    
    for (x, y) in points:
        # 计算特征图上的有效粘贴边界，避免越界
        x_min = max(0, int(x - kernel_radius))
        x_max = min(W, int(x + kernel_radius + 1))
        y_min = max(0, int(y - kernel_radius))
        y_max = min(H, int(y + kernel_radius + 1))
        
        # 计算高斯核的裁剪区域，适配边缘坐标
        k_x_min = max(0, kernel_radius - int(x))
        k_x_max = kernel_size - max(0, int(x + kernel_radius + 1) - W)
        k_y_min = max(0, kernel_radius - int(y))
        k_y_max = kernel_size - max(0, int(y + kernel_radius + 1) - H)
        
        # 粘贴到特征图，重叠区域累加
        feat_map[y_min:y_max, x_min:x_max] += gaussian_kernel[k_y_min:k_y_max, k_x_min:k_x_max]
    
    return feat_map

3. 大量坐标点场景：卷积实现 #

适合坐标点数量>1000的场景，用卷积替代循环提升效率：

def batch_draw_gaussian_conv(points, H, W, kernel_radius=3, sigma=None):
    kernel_size = 2 * kernel_radius + 1
    gaussian_kernel = generate_2d_gaussian_kernel(kernel_size, sigma).to(points.device)
    # 初始化点掩码图，对应坐标设为1
    point_map = torch.zeros((1, 1, H, W), device=points.device)
    for (x, y) in points:
        point_map[0, 0, int(y), int(x)] = 1.0
    # 高斯卷积，padding保证输出尺寸和输入一致
    feat_map = F.conv2d(point_map, gaussian_kernel.unsqueeze(0).unsqueeze(0), padding=kernel_radius)
    return feat_map.squeeze()

使用示例 #

# 测试：3个坐标点生成64x64的特征图
points = torch.tensor([[10, 20], [30, 15], [50, 50]])
feat_map = draw_gaussian_on_featuremap(points, H=64, W=64, kernel_radius=4)

注：如果输入为浮点型亚像素坐标，可先对坐标做四舍五入取整，或基于坐标的小数部分做双线性插值加权粘贴高斯核，上述代码默认适配整数坐标场景。

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