如何构建一个 3D 张量,其中每个 2D 子张量都是 PyTorch 中的对角矩阵?

How to construct a 3D Tensor where every 2D sub tensor is a diagonal matrix in PyTorch?

假设我有二维张量,index_in_batch * diag_ele。 如何得到一个3D张量index_in_batch * Matrix(谁是对角矩阵,由drag_ele构造)?

torch.diag()只在输入为一维时构造对角矩阵,输入为二维时return构造对角矩阵

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]

自动后退的解决方法,包裹在Variable中

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)

使用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]]])