python - 在 PyTorch 中加速 SVD

Question

我正在使用 Pytorch 为 CIFAR10 执行一些分类任务，对于每次迭代，我都必须对每个批处理进行一些预处理，然后才能将其反馈给模型。 以下是每个批处理的预处理部分的代码:

S = torch.zeros((batch_size, C, H, W))
for i in range(batch_size):
    img = batch[i, :, :, :]
    for c in range(C):                
        U, _, V = torch.svd(img[c])
        S[i, c] = U[:, 0].view(-1, 1).matmul(V[:, 0].view(1, -1))

但是，这个计算速度非常慢。有什么方法可以加快这段代码的速度吗？

Answer 1

批量计算

假设您有 PyTorch >= 1.2.0，则支持批量 SVD，因此您可以使用

U, _, V = torch.svd(batch)
S = U[:, :, :, 0].unsqueeze(3) @ V[:, :, :, 0].unsqueeze(2)

我发现它比迭代版本平均要快一点。

<小时/>

截断的 SVD(仅限 CPU)

如果您没有 cuda 加速，您可以使用截断的 SVD 来避免计算不必要的奇异值/向量。不幸的是 PyTorch 不支持截断的 SVD 并且据我所知没有可用的批处理或 GPU 版本。我知道有两种选择

• scipy.sparse.linalg.svds
• sklearn.sparse.linalg.randomized_svd

这两个选项都允许您选择要返回的组件数量。在OP的原始问题中，我们只想要第一个组件。

尽管我没有在稀疏矩阵上使用它，但我发现使用 k=1 的 svds 比 torch.svd 快 10 倍左右> 关于 CPU 张量。我发现 randomized_svd 仅快 2 倍左右。您的结果将取决于实际数据。此外，svds 应该比 randomized_svd 更准确一些。请记住，这些结果与 torch.svd 结果之间会有微小差异，但它们应该可以忽略不计。

import scipy.sparse.linalg as sp
import numpy as np

S = torch.zeros((batch_size, C, H, W))
for i in range(batch_size):
    img = batch[i, :, :, :]
    for c in range(C):
        u, _, v = sp.svds(img[c], k=1)
        S[i, c] = torch.from_numpy(np.outer(u, v))

<小时/>