用torch写高斯核函数

下面是用 PyTorch 编写高斯核函数的代码：

import torch

def gaussian_kernel(size: int, sigma: float) -> torch.Tensor:
    # Create a 1D Gaussian kernel
    kernel_1d = torch.exp(
        -torch.pow(torch.arange(size).float() - size // 2, 2) / (2 * sigma ** 2))
    # Normalize the kernel
    kernel_1d /= torch.sum(kernel_1d)
    # Convert to a 2D kernel
    kernel_2d = torch.outer(kernel_1d, kernel_1d)
    return kernel_2d

其中，`size` 是指高斯核的大小（即宽度和高度），`sigma` 是高斯分布的标准差。该函数会返回一个形状为 `(size, size)` 的 2D 张量，表示高斯核矩阵。

在该函数中，我们首先创建一个 1D 的高斯核，其大小为 `size`，然后将其规范化，使得所有元素的和为 1。接着，我们使用 `torch.outer` 函数将 1D 高斯核转换为 2D 高斯核。具体来说，该函数会将两个向量张量的外积计算出来，如下所示：

torch.outer(torch.tensor([1, 2, 3]), torch.tensor([4, 5, 6]))
# Output: tensor([[ 4,  5,  6],
#                 [ 8, 10, 12],
#                 [12, 15, 18]])

最后，我们返回 2D 高斯核矩阵。