Python PyTorch solve用法及代码示例

本文简要介绍python语言中 torch.linalg.solve 的用法。

用法:

torch.linalg.solve(A, B, *, out=None) → Tensor

参数：

• A(Tensor) -形状为 (*, n, n) 的张量，其中 * 是零个或多个批次维度。

• B(Tensor) -根据上述规则，形状为 (*, n) 或 (*, n, k) 或 (n,) 或 (n, k) 的右侧张量

关键字参数：

out(Tensor,可选的) -输出张量。如果 None 则忽略。默认值：None。

抛出：

RuntimeError - 如果 A 矩阵不可逆或批处理中的任何矩阵 A 不可逆。

计算具有唯一解的线性方程组的解。

让\mathbb{K}是\mathbb{R}或者\mathbb{C}, 这个函数计算解X \in \mathbb{K}^{n \times k}的线性系统关联到A \in \mathbb{K}^{n \times n}, B \in \mathbb{K}^{m \times k}, 定义为

AX = B

当且仅当 A 为 invertible 时，此线性方程组有一个解。此函数假定A 是可逆的。

支持 float、double、cfloat 和 cdouble dtypes 的输入。还支持矩阵批次，如果输入是矩阵批次，则输出具有相同的批次尺寸。

让* 是零个或多个批次维度，

• 如果 A 的形状为 (*, n, n) 且 B 的形状为 (*, n)(一组向量)或 (*, n, k) 形状(一组矩阵或“多个右侧”)，则此函数返回 X形状分别为 (*, n) 或 (*, n, k)。

• 否则，如果 A 具有形状 (*, n, n) 并且 B 具有形状 (n,) 或 (n, k) ，则广播 B 分别具有形状 (*, n) 或 (*, n, k) 。然后，此函数返回生成的一批线性方程组的解。

注意

此函数计算 X = A .inverse() @ B 比单独执行计算更快且数值更稳定。

注意

当输入在 CUDA 设备上时，此函数将该设备与 CPU 同步。

例子：

>>> A = torch.randn(3, 3)
>>> b = torch.randn(3)
>>> x = torch.linalg.solve(A, b)
>>> torch.allclose(A @ x, b)
True
>>> A = torch.randn(2, 3, 3)
>>> B = torch.randn(2, 3, 4)
>>> X = torch.linalg.solve(A, B)
>>> X.shape
torch.Size([2, 3, 4])
>>> torch.allclose(A @ X, B)
True

>>> A = torch.randn(2, 3, 3)
>>> b = torch.randn(3, 1)
>>> x = torch.linalg.solve(A, b) # b is broadcasted to size (2, 3, 1)
>>> x.shape
torch.Size([2, 3, 1])
>>> torch.allclose(A @ x, b)
True
>>> b = torch.randn(3)
>>> x = torch.linalg.solve(A, b) # b is broadcasted to size (2, 3)
>>> x.shape
torch.Size([2, 3])
>>> Ax = A @ x.unsqueeze(-1)
>>> torch.allclose(Ax, b.unsqueeze(-1).expand_as(Ax))
True