torch 如何生成主对角阵?

高山莫衣2023-12-12 20:39

可以使用 torch.eye() 函数生成主对角阵，并将对角元素设置为 0 或 1。具体来说，如果需要生成对角元素都为 0 的主对角阵，可以直接调用 torch.eye() 函数，并指定对角线上的元素为 0。如果需要生成对角元素都为 1 的主对角阵，则可以先生成对角线为 1 的主对角阵，然后将对角线上的元素全部乘以 0。

下面是生成上述两种主对角阵的示例代码：

python 复制代码

import torch
# 生成对角元素都为 0 的主对角阵
n = 5
A = torch.eye(n) * 0
print(A)

输出结果为：

python 复制代码

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

python 复制代码

import torch
# 生成对角元素都为 1 的主对角阵
n = 5
B = torch.eye(n)
print(B)

输出结果为：

python 复制代码

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

除了使用 torch.eye() 函数生成主对角阵之外，还可以使用 torch.diag() 函数来生成主对角阵。torch.diag() 函数既可以将一个向量转化为对角矩阵，也可以将一个对角矩阵中的对角元素抽取出来形成一个向量。

下面是使用 torch.diag() 函数生成主对角阵的示例代码：

python 复制代码

import torch
# 生成对角元素都为 0 的主对角阵
n = 5
A = torch.zeros(n, n)
A = torch.diag(torch.zeros(n))
print(A)

输出结果为：

python 复制代码

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

python 复制代码

import torch
# 生成对角元素都为 1 的主对角阵
n = 5
B = torch.ones(n, n)
B = torch.diag(torch.ones(n))
print(B)

输出结果为：

python 复制代码

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