Pytorch：将两个高维张量 (2, 5, 3) * (2, 5) 乘以 (2, 5, 3)

Question

我想将两个高维张量 (2, 5, 3) * (2, 5) 乘以 (2, 5, 3)，将每个行向量乘以一个标量。

例如

emb = nn.Embedding(6, 3)

input = torch.tensor([[1, 2, 3, 4, 5,],
                      [2, 3, 1, 4, 5,]])
input_emb = emb(input)


print(input.shape)
> torch.Size([2, 5])

print(input_emb.shape)
> torch.Size([2, 5, 3])

print(input_emb)
> tensor([[[-1.9114, -0.1580,  1.2186],
         [ 0.4627,  0.9119, -1.1691],
         [ 0.6452, -0.6944,  1.9659],
         [-0.5048,  0.6411, -1.3568],
         [-0.2328, -0.9498,  0.7216]],

        [[ 0.4627,  0.9119, -1.1691],
         [ 0.6452, -0.6944,  1.9659],
         [-1.9114, -0.1580,  1.2186],
         [-0.5048,  0.6411, -1.3568],
         [-0.2328, -0.9498,  0.7216]]], grad_fn=<EmbeddingBackward>)

我想乘法如下：

// It is written in this way for convenience, not mathematical true. 

// multiply each row vector by a scalar
[[
         [-1.9114, -0.1580,  1.2186] * 1
         [ 0.4627,  0.9119, -1.1691] * 2
         [ 0.6452, -0.6944,  1.9659] * 3
         [-0.5048,  0.6411, -1.3568] * 4
         [-0.2328, -0.9498,  0.7216] * 5
] 
[
         [ 0.4627,  0.9119, -1.1691] * 2
         [ 0.6452, -0.6944,  1.9659] * 3
         [-1.9114, -0.1580,  1.2186] * 1
         [-0.5048,  0.6411, -1.3568] * 4
         [-0.2328, -0.9498,  0.7216] * 5
]]

除了多for循环方式，如何用PyTorch个API简洁的实现？
提前致谢。

Answer 1

您可以通过正确对齐两个张量的维度：

import torch
from torch.nn import Embedding

emb = Embedding(6, 3)
inp = torch.tensor([[1, 2, 3, 4, 5,],
                      [2, 3, 1, 4, 5,]])
input_emb = emb(inp)

inp[...,None] * input_emb

tensor([[[-0.3069, -0.7727, -0.3772],
         [-2.8308,  1.3438, -1.1167],
         [ 0.6366,  0.6509, -3.2282],
         [-4.3004,  3.2342, -0.6556],
         [-3.0045, -0.0191, -7.4436]],

        [[-2.8308,  1.3438, -1.1167],
         [ 0.6366,  0.6509, -3.2282],
         [-0.3069, -0.7727, -0.3772],
         [-4.3004,  3.2342, -0.6556],
         [-3.0045, -0.0191, -7.4436]]], grad_fn=<MulBackward0>)

Pytorch：将两个高维张量 (2, 5, 3) * (2, 5) 乘以 (2, 5, 3)

Pytorch: multiply two high dimensions tensor, (2, 5, 3) * (2, 5) into (2, 5, 3)

python

pytorch

tensor