pytorch将向量拉成对角阵（计算图的度矩阵）

在使用pytorch编写神经网络时，可能会用到创建对角阵（例如图邻接矩阵的度矩阵）。最简单的方法是创建一个随机的矩阵然后循环赋值，但在数据量较大时可能会比较慢，此时可以先获得对角阵对角线上的值构成的向量，然后将该向量拉成对角阵。例如，有一邻接矩阵A，现在要求其度矩阵D（对角阵，每一行的对角元是A中对应节点的度），举例如下：

import torch

A=torch.ones(3,3)
print(A)
d=torch.sum(A,dim=1)
print(d)
D=torch.diag(d)
print(D)

对应的输出

tensor([[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]])
tensor([3., 3., 3.])
tensor([[3., 0., 0.],
        [0., 3., 0.],
        [0., 0., 3.]])