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