如何构造一个三维张量,其中每个二维子张量是PyTorch中的对角矩阵?

import torch

a = torch.rand(2, 3)
print(a)
b = torch.eye(a.size(1))
c = a.unsqueeze(2).expand(*a.size(), a.size(1))
d = c * b
print(d)

输出

 0.5938  0.5769  0.0555
 0.9629  0.5343  0.2576
[torch.FloatTensor of size 2x3]


(0 ,.,.) = 
  0.5938  0.0000  0.0000
  0.0000  0.5769  0.0000
  0.0000  0.0000  0.0555

(1 ,.,.) = 
  0.9629  0.0000  0.0000
  0.0000  0.5343  0.0000
  0.0000  0.0000  0.2576
[torch.FloatTensor of size 2x3x3]

使用 torch.diag_embed :

>>> a = torch.randn(2, 3)
>>> torch.diag_embed(a)
tensor([[[ 1.5410,  0.0000,  0.0000],
         [ 0.0000, -0.2934,  0.0000],
         [ 0.0000,  0.0000, -2.1788]],

        [[ 0.5684,  0.0000,  0.0000],
         [ 0.0000, -1.0845,  0.0000],
         [ 0.0000,  0.0000, -1.3986]]])

通过在变量中包装实现自动向后的解决方案。

import torch

a = torch.rand(2, 3)
print(a)

b = Variable(torch.eye(a.size(1)))
c = a.unsqueeze(2).expand(*a.size(), degree_inv.size(1))
b_expand =  b.unsqueeze(0).expand(c.size(0), *b.size())
d = torch.mul(c.double(), b_expand.double())

print(d)