PyTorch torch.linalg.lstsq 函数
torch.linalg.lstsq 是 PyTorch 线性代数模块中用于求解线性最小二乘问题的函数。它求解 argmin_x ||AX - B||_F。
函数定义
torch.linalg.lstsq(A, B, rcond=None, driver=None)
参数:
A(Tensor): 系数矩阵。B(Tensor): 右侧矩阵或向量。rcond(float, 可选): 用于截断奇异值的条件数。driver(str, 可选): 求解器选择,可选 'gels', 'gelsd', 'gelsy'。
返回值:
torch.Tensor: 返回最小二乘解。
使用示例
实例
import torch
# 创建系数矩阵和右侧向量
A = torch.tensor([[1.0, 1.0], [1.0, 2.0], [1.0, 3.0]], dtype=torch.float64)
B = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64)
# 最小二乘求解
X = torch.linalg.lstsq(A, B).solution
print("系数矩阵 A:")
print(A)
print("n右侧向量 B:")
print(B)
print("n最小二乘解 X:")
print(X)
# 创建系数矩阵和右侧向量
A = torch.tensor([[1.0, 1.0], [1.0, 2.0], [1.0, 3.0]], dtype=torch.float64)
B = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64)
# 最小二乘求解
X = torch.linalg.lstsq(A, B).solution
print("系数矩阵 A:")
print(A)
print("n右侧向量 B:")
print(B)
print("n最小二乘解 X:")
print(X)
输出结果为:
系数矩阵 A:
tensor([[1., 1.],
[1., 2.],
[1., 3.]], dtype=torch.float64)
右侧向量 B:
tensor([1., 2., 3.], dtype=float64)
最小二乘解 X:
tensor([0.0000, 1.0000], dtype=float64)

Pytorch torch 参考手册