mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
2ec779d46c
Summary: This PR implements the following linear algebra algorithms for low-rank matrices: - [x] Approximate `A` as `Q Q^H A` - using Algorithm 4.4 from [Halko et al, 2009](http://arxiv.org/abs/0909.4061). + exposed as `torch.lowrank.get_approximate_basis(A, q, niter=2, M=None) -> Q` + [x] dense matrices + [x] batches of dense matrices + [x] sparse matrices + [x] documentation - [x] SVD - using Algorithm 5.1 from [Halko et al, 2009](http://arxiv.org/abs/0909.4061). + uses `torch.lowrank.get_approximate_basis` + exposed as `torch.svd_lowrank(A, q=6, niter=2, M=None) -> (U, S, V)` + [x] dense matrices + [x] batches of dense matrices + [x] sparse matrices + [x] documentation - [x] PCA - using `torch.svd_lowrank` + uses `torch.svd_lowrank` + exposed as `torch.pca_lowrank(A, center=True, q=None, niter=2) -> (U, S, V)` + [x] dense matrices + [x] batches of dense matrices + [x] sparse matrices, uses non-centered sparse matrix algorithm + [x] documentation - [x] generalized eigenvalue solver using the original LOBPCG algorithm [Knyazev, 2001](https://epubs.siam.org/doi/abs/10.1137/S1064827500366124) + exposed as `torch.lobpcg(A, B=None, k=1, method="basic", ...)` + [x] dense matrices + [x] batches of dense matrices + [x] sparse matrices + [x] documentation - [x] generalized eigenvalue solver using robust LOBPCG with orthogonal basis selection [Stathopoulos, 2002](https://epubs.siam.org/doi/10.1137/S1064827500370883) + exposed as `torch.lobpcg(A, B=None, k=1, method="ortho", ...)` + [x] dense matrices + [x] batches of dense matrices + [x] sparse matrices + [x] documentation - [x] generalized eigenvalue solver using the robust and efficient LOBPCG Algorithm 8 from [Duersch et al, 2018](https://epubs.siam.org/doi/abs/10.1137/17M1129830) that switches to orthogonal basis selection automatically + the "ortho" method improves iterations so rapidly that in the current test cases it does not make sense to use the basic iterations at all. If users will have matrices for which basic iterations could improve convergence then the `tracker` argument allows breaking the iteration process at user choice so that the user can switch to the orthogonal basis selection if needed. In conclusion, there is no need to implement Algorithm 8 at this point. - [x] benchmarks + [x] `torch.svd` vs `torch.svd_lowrank`, see notebook [Low-rank SVD](https://github.com/Quansight/pearu-sandbox/blob/master/pytorch/Low-rank%20SVD.ipynb). In conclusion, the low-rank SVD is going to be useful only for large sparse matrices where the full-rank SVD will fail due to memory limitations. + [x] `torch.lobpcg` vs `scipy.sparse.linalg.lobpcg`, see notebook [LOBPCG - pytorch vs scipy](https://github.com/Quansight/pearu-sandbox/blob/master/pytorch/LOBPCG%20-%20pytorch%20vs%20scipy.ipynb). In conculsion, both implementations give the same results (up to numerical errors from different methods), scipy lobpcg implementation is generally faster. + [x] On very small tolerance cases, `torch.lobpcg` is more robust than `scipy.sparse.linalg.lobpcg` (see `test_lobpcg_scipy` results) Resolves https://github.com/pytorch/pytorch/issues/8049. Pull Request resolved: https://github.com/pytorch/pytorch/pull/29488 Differential Revision: D20193196 Pulled By: vincentqb fbshipit-source-id: 78a4879912424595e6ea95a95e483a37487a907e
102 lines
2.3 KiB
Python
102 lines
2.3 KiB
Python
"""Various linear algebra utility methods for internal use.
|
|
|
|
"""
|
|
|
|
import torch
|
|
|
|
|
|
def is_sparse(A):
|
|
"""Check if tensor A is a sparse tensor"""
|
|
if isinstance(A, torch.Tensor):
|
|
return A.layout == torch.sparse_coo
|
|
raise TypeError("expected Tensor but got %s" % (type(A).__name__))
|
|
|
|
|
|
def get_floating_dtype(A):
|
|
"""Return the floating point dtype of tensor A.
|
|
|
|
Integer types map to float32.
|
|
"""
|
|
dtype = A.dtype
|
|
if dtype in (torch.float16, torch.float32, torch.float64):
|
|
return dtype
|
|
return torch.float32
|
|
|
|
|
|
def matmul(A, B):
|
|
# type: (Optional[Tensor], Tensor) -> Tensor
|
|
"""Multiply two matrices.
|
|
|
|
If A is None, return B. A can be sparse or dense. B is always
|
|
dense.
|
|
"""
|
|
if A is None:
|
|
return B
|
|
if is_sparse(A):
|
|
return torch.sparse.mm(A, B)
|
|
return torch.matmul(A, B)
|
|
|
|
|
|
def conjugate(A):
|
|
"""Return conjugate of tensor A.
|
|
|
|
.. note:: If A's dtype is not complex, A is returned.
|
|
"""
|
|
if A.is_complex():
|
|
return A.conj()
|
|
return A
|
|
|
|
|
|
def transpose(A):
|
|
"""Return transpose of a matrix or batches of matrices.
|
|
"""
|
|
ndim = len(A.shape)
|
|
return A.transpose(ndim - 1, ndim - 2)
|
|
|
|
|
|
def transjugate(A):
|
|
"""Return transpose conjugate of a matrix or batches of matrices.
|
|
"""
|
|
return conjugate(transpose(A))
|
|
|
|
|
|
def bform(X, A, Y):
|
|
# type: (Tensor, Optional[Tensor], Tensor) -> Tensor
|
|
"""Return bilinear form of matrices: :math:`X^T A Y`.
|
|
"""
|
|
return matmul(transpose(X), matmul(A, Y))
|
|
|
|
|
|
def qform(A, S):
|
|
# type: (Optional[Tensor], Tensor) -> Tensor
|
|
"""Return quadratic form :math:`S^T A S`.
|
|
"""
|
|
return bform(S, A, S)
|
|
|
|
|
|
def basis(A):
|
|
"""Return orthogonal basis of A columns.
|
|
"""
|
|
if A.is_cuda:
|
|
# torch.orgqr is not available in CUDA
|
|
Q, _ = torch.qr(A, some=True)
|
|
else:
|
|
Q = torch.orgqr(*torch.geqrf(A))
|
|
return Q
|
|
|
|
|
|
def symeig(A, largest=False, eigenvectors=True):
|
|
# type: (Tensor, Optional[bool], Optional[bool]) -> Tuple[Tensor, Tensor]
|
|
"""Return eigenpairs of A with specified ordering.
|
|
"""
|
|
if largest is None:
|
|
largest = False
|
|
if eigenvectors is None:
|
|
eigenvectors = True
|
|
E, Z = torch.symeig(A, eigenvectors, True)
|
|
# assuming that E is ordered
|
|
if largest:
|
|
E = torch.flip(E, dims=(-1,))
|
|
Z = torch.flip(Z, dims=(-1,))
|
|
return E, Z
|