PyTorch torch.trace 函数
Pytorch torch 参考手册torch.trace 是 PyTorch 中用于计算矩阵迹的函数。迹是矩阵主对角线元素的和。
函数定义
torch.trace(input)
参数说明：
input: 输入张量（至少是二维的）
使用示例
实例
import torch
# 创建矩阵
a = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
# 计算迹（主对角线元素之和）
y = torch.trace(a)
print(y)
# 创建矩阵
a = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
# 计算迹（主对角线元素之和）
y = torch.trace(a)
print(y)
输出结果为：
tensor(15)
实例
import torch
# 创建非方阵
a = torch.tensor([[1, 2, 3], [4, 5, 6]])
# 计算迹
y = torch.trace(a)
print(y)
# 创建非方阵
a = torch.tensor([[1, 2, 3], [4, 5, 6]])
# 计算迹
y = torch.trace(a)
print(y)
输出结果为：
tensor(6)
实例
import torch
# 创建多维张量，只考虑前两个维度
a = torch.randn(3, 4, 4, 5)
y = torch.trace(a)
print(y.shape)
# 创建多维张量，只考虑前两个维度
a = torch.randn(3, 4, 4, 5)
y = torch.trace(a)
print(y.shape)
输出结果为：
torch.Size([3, 5])