torch.nn.modules.distance 的源代码

from .module import Module
from .. import functional as F

from torch import Tensor

__all__ = ['PairwiseDistance', 'CosineSimilarity']

class PairwiseDistance(Module):
    r"""
    计算输入向量之间的成对距离，或输入矩阵的列之间的成对距离。

    距离使用 ``p``-范数计算，并添加常数 ``eps`` 以避免除以零
    如果 ``p`` 为负数，即：

    .. math ::
        \mathrm{dist}\left(x, y\right) = \left\Vert x-y + \epsilon e \right\Vert_p,

    其中 :math:`e` 是全一向量，``p``-范数由下式给出。

    .. math ::
        \Vert x \Vert _p = \left( \sum_{i=1}^n  \vert x_i \vert ^ p \right) ^ {1/p}.

    参数:
        p (实数, 可选): 范数度。可以是负数。默认值: 2
        eps (浮点数, 可选): 用于避免除以零的小值。
            默认值: 1e-6
        keepdim (布尔值, 可选): 确定是否保留向量维度。
            默认值: False
    形状:
        - 输入1: :math:`(N, D)` 或 :math:`(D)` 其中 `N = 批次维度` 和 `D = 向量维度`
        - 输入2: :math:`(N, D)` 或 :math:`(D)`，与输入1形状相同
        - 输出: :math:`(N)` 或 :math:`()` 基于输入维度。
          如果 :attr:`keepdim` 为 ``True``，则 :math:`(N, 1)` 或 :math:`(1)` 基于输入维度。

    示例::
        >>> pdist = nn.PairwiseDistance(p=2)
        >>> input1 = torch.randn(100, 128)
        >>> input2 = torch.randn(100, 128)
        >>> output = pdist(input1, input2)
    """

    __constants__ = ['norm', 'eps', 'keepdim']
    norm: float
    eps: float
    keepdim: bool

    def __init__(self, p: float = 2., eps: float = 1e-6, keepdim: bool = False) -> None:
        super().__init__()
        self.norm = p
        self.eps = eps
        self.keepdim = keepdim

    def forward(self, x1: Tensor, x2: Tensor) -> Tensor:
        return F.pairwise_distance(x1, x2, self.norm, self.eps, self.keepdim)


class CosineSimilarity(Module):
    r"""返回 :math:`x_1` 和 :math:`x_2` 之间的余弦相似度，沿 `dim` 计算。

    .. math ::
        \text{similarity} = \dfrac{x_1 \cdot x_2}{\max(\Vert x_1 \Vert _2 \cdot \Vert x_2 \Vert _2, \epsilon)}.

    参数:
        dim (整数, 可选): 计算余弦相似度的维度。默认值: 1
        eps (浮点数, 可选): 用于避免除以零的小值。
            默认值: 1e-8
    形状:
        - 输入1: :math:`(\ast_1, D, \ast_2)` 其中 D 位于 `dim` 位置
        - 输入2: :math:`(\ast_1, D, \ast_2)`，与 x1 的维度数相同，在 `dim` 维度上与 x1 匹配，
              并且在其他维度上可与 x1 广播。
        - 输出: :math:`(\ast_1, \ast_2)`
    示例::
        >>> input1 = torch.randn(100, 128)
        >>> input2 = torch.randn(100, 128)
        >>> cos = nn.CosineSimilarity(dim=1, eps=1e-6)
        >>> output = cos(input1, input2)
    """

    __constants__ = ['dim', 'eps']
    dim: int
    eps: float

    def __init__(self, dim: int = 1, eps: float = 1e-8) -> None:
        super().__init__()
        self.dim = dim
        self.eps = eps

    def forward(self, x1: Tensor, x2: Tensor) -> Tensor:
        return F.cosine_similarity(x1, x2, self.dim, self.eps)