It is a batch version SVD of pytorch implemented by cuSolver including forward and backward function.
def batch_svd(x):
"""
input:
x --- shape of [B, M, N]
return:
U, S, V = batch_svd(x) where x = USV^T
"""
The forward function is modified from ShigekiKarita/pytorch-cusolver and I fixed several bugs of it. The backward function is adapted from pytorch official svd backward function. I converted it to a batch version.
NOTE: batch_svd
only supports CudaFloatTensor
now. Other types may be supported in the future.
-
Pytorch >= 1.0
diag_embed() is used in torch_batch_svd.cpp at the backward function. Pytorch with version lower than 1.0 does not contains diag_embed(). If you want to use it in lower version pytorch, you can replace diag_embed() by some existing function.
-
CUDA 9.0
export CUDA_HOME=/your/cuda/home/directory/
export LIBRARY_PATH=$LIBRARY_PATH:/your/cuda/lib64/ (optional)
python setup.py install
python test.py
-
batch_svd()
has no configurations ofsome
,compute_uv
liketorch.svd()
.batch_svd(x)
is equivalent totorch.svd(x, some=True, compute_uv=True)
. -
The sign of column vectors at U and V may be different from
torch.svd()
. -
batch_svd()
is much more faster thantorch.svd()
using loop.
See test.py
.