Home

Awesome

Pytorch Batched SVD

Introduction

A 100x faster SVD for PyTorch including forward and backward function.

Performance:

matrix sizetorch_batch_svd.svdtorch.svd
(10000, 9, 9)0.043 s19.352 s
(20000, 9, 9)0.073 s34.578 s
import torch
from torch_batch_svd import svd

A = torch.rand(1000000, 3, 3).cuda()
u, s, v = svd(A)
u, s, v = torch.svd(A)  # probably you should take a coffee break here

The catch here is that it only works for matrices whose row and column are smaller than 32. Other than that, torch_batch_svd.svd can be a drop-in for the native one.

The forward function is modified from ShigekiKarita/pytorch-cusolver and I fixed several bugs of it. The backward function is borrowed from the PyTorch official svd backward function. I converted it to a batched version.

NOTE: batch_svd supports all torch.half, torch.float and torch.double tensors now.

NOTE: SVD for torch.half is performed by casting to torch.float as there is no CuSolver implementation for c10::half.

NOTE: Sometimes, tests will fail for torch.double tensor due to numerical imprecision.

Get Started

Requirements

Install

export CUDA_HOME=/your/cuda/home/directory/
python setup.py install

Test

cd tests
python test.py

Differences between torch.svd()

Example

See test.py and introduction.

Star History

Star History Chart