当前位置: 首页>>代码示例 >>用法及示例精选 >>正文


Python PyTorch tensordot用法及代码示例


本文简要介绍python语言中 torch.tensordot 的用法。

用法:

torch.tensordot(a, b, dims=2, out=None)

参数

  • a(Tensor) -左张量收缩

  • b(Tensor) -右张量收缩

  • dims(int或者元组[List[int],List[int]] 或者List[List[int]]包含两个列表或者Tensor) - 要收缩的维度数量或明确的维度列表ab分别

返回 a 和 b 在多个维度上的收缩。

tensordot 实现广义矩阵乘积。

当使用非负整数参数 dims = 调用,并且 ab 的维数分别为 时,tensordot() 计算

当使用列表形式的dims调用时,给定的维度将被收缩以代替a的最后一个 的第一个 。这些维度中的大小必须匹配,但 tensordot() 将处理广播维度。

例子:

>>> a = torch.arange(60.).reshape(3, 4, 5)
>>> b = torch.arange(24.).reshape(4, 3, 2)
>>> torch.tensordot(a, b, dims=([1, 0], [0, 1]))
tensor([[4400., 4730.],
        [4532., 4874.],
        [4664., 5018.],
        [4796., 5162.],
        [4928., 5306.]])

>>> a = torch.randn(3, 4, 5, device='cuda')
>>> b = torch.randn(4, 5, 6, device='cuda')
>>> c = torch.tensordot(a, b, dims=2).cpu()
tensor([[ 8.3504, -2.5436,  6.2922,  2.7556, -1.0732,  3.2741],
        [ 3.3161,  0.0704,  5.0187, -0.4079, -4.3126,  4.8744],
        [ 0.8223,  3.9445,  3.2168, -0.2400,  3.4117,  1.7780]])

>>> a = torch.randn(3, 5, 4, 6)
>>> b = torch.randn(6, 4, 5, 3)
>>> torch.tensordot(a, b, dims=([2, 1, 3], [1, 2, 0]))
tensor([[  7.7193,  -2.4867, -10.3204],
        [  1.5513, -14.4737,  -6.5113],
        [ -0.2850,   4.2573,  -3.5997]])

相关用法


注:本文由纯净天空筛选整理自pytorch.org大神的英文原创作品 torch.tensordot。非经特殊声明,原始代码版权归原作者所有,本译文未经允许或授权,请勿转载或复制。