如何批量把tensor矩阵的对角置零

可以利用torch.diag()函数来获取矩阵的对角线，然后将其置零，最后再将对角线重新填回原矩阵中。以下是示例代码：

import torch

# 创建一个5x5的随机张量
tensor = torch.randn(5, 5)

# 获取对角线
diag = torch.diag(tensor)

# 将对角线置零
diag.fill_(0)

# 将修改后的对角线重新填回原矩阵中
tensor = torch.diag_embed(diag)

# 批量将矩阵的对角线置零
def batch_diag_zero(tensor_list):
    for i in range(len(tensor_list)):
        diag = torch.diag(tensor_list[i])
        diag.fill_(0)
        tensor_list[i] = torch.diag_embed(diag)

# 示例
tensor_list = [torch.randn(3, 3), torch.randn(4, 4)]
batch_diag_zero(tensor_list)
print(tensor_list)

其中，torch.diag_embed()函数可以将对角线转换为对角矩阵。

相关问题

但我的tensor为(16,14,14),我要把每个（14,14）的矩阵的对角置零

可以利用PyTorch的高级索引功能来批量对每个（14,14）的小矩阵进行操作，具体代码如下：

import torch

# 创建一个形状为(16, 14, 14)的随机张量
tensor = torch.randn(16, 14, 14)

# 获取每个小矩阵的对角线并置零
idx = torch.arange(14)
tensor[:, idx, idx] = 0

# 查看结果
print(tensor)

其中，`torch.arange(14)`用于创建一个长度为14的序列，表示矩阵的行列索引。通过高级索引`tensor[:, idx, idx]`可以获取每个小矩阵的对角线，并将其置零。