RAM 在计算上爆炸

Question

下面是使用dask 和cupy 的可运行代码sn-p，我遇到了问题。我在激活了 GPU 的 Google Colab 上运行它。

基本上我的问题是，A 和 At 是对 RAM 来说太大的数组，这就是我使用 Dask 的原因。在这些对于 RAM 数组来说太大的情况下，我运行操作，但我想获得 AtW1[:,k] （作为一个cupy数组）而不会耗尽我的 RAM 或 GPU 内存，因为我需要这个值用于进一步的操作。我怎样才能做到这一点？

import dask.array as da
import cupy as cp

from dask_cuda import LocalCUDACluster
from dask.distributed import Client

cluster = LocalCUDACluster()
client = Client(cluster)

n_topics = 10
n_terms = 10000
n_docs = 250000
k = 0

A = da.zeros_like(cp.zeros(1), shape=(n_terms,n_docs), chunks=(1,n_docs))
At = A.transpose().rechunk((10, n_terms))

n_row = A.shape[0]
n_col = A.shape[1]

W1 = cp.random.random((n_row, n_topics))
H = cp.random.random((n_col, n_topics))
W1tW1 = W1.T.dot(W1)

# The problem starts here: i want to get AtW1[:,k] as a cupy array without blowing my RAM
AtW1 = da.dot(At, W1).rechunk((n_docs, 1))
AtW1 = AtW1.persist()
val = AtW1[:,k].compute()
val

编辑：另一个有同样问题的例子。如何在不更改第一部分的情况下实现计算？第一部分的变量就在那里，所以数组的维度是明确的，例子是可运行的。

import cupy as cp
import dask.array as da
from dask.array.linalg import norm
from dask_cuda import LocalCUDACluster
from dask.distributed import Client

cluster = LocalCUDACluster()
client = Client(cluster)

####### These are given and cannot be changed #######
big = 250000
small = 10000
full = da.full_like(cp.full(1, 2.0), fill_value=2.0, shape=(small,big), chunks=(1,big))
a = cp.random.random((small, 10))
b = cp.transpose(cp.random.random((big, 10)))

###### How can i make this work? #######
ab = da.dot(da.from_array(a), da.from_array(b))
subtracted = full - ab
normalized = norm(subtracted, ord='fro')
val = normalized.compute()

Answer 1

虽然重新分块的想法在纸面上很有意义，但在实践中重新分块需要非常小心，因为它只能重塑原则上可以被阻止的工作。

例如，比较以下两种方法：

A = da.zeros_like(cp.zeros(1), shape=(n_terms,n_docs), chunks=(1,n_docs))
At = A.transpose().rechunk((1, n_terms))

At_version2 = da.zeros_like(cp.zeros(1), shape=(n_docs, n_terms), chunks=(1,n_terms))

当请求At[0,:] 时，dask 将完成的工作量将远远大于为 At_version2[0,:] 完成的工作量。为什么？因为At 的原始数据将在(1,n_docs) 的块中定义，所以要生成转置的第一行，dask 必须在每一行计算原始数组A。

使用此逻辑，最好确保您的数据以确实与获得最终结果所需的块乘法兼容的形状提供。这将导致更少的任务和更低的内存需求，因为现在只需将数据的子集保存在内存中。

这是快速计算的代码的修改版本：

import dask.array as da
import cupy as cp

from dask_cuda import LocalCUDACluster
from dask.distributed import Client

cluster = LocalCUDACluster()
client = Client(cluster)

n_topics = 10
n_terms = 10000
n_docs = 250000
k = 0

At_v2 = da.zeros_like(cp.zeros(1), shape=(n_docs, n_terms), chunks=(1000,n_terms))

W1 = cp.random.random((n_terms, n_topics))

result = da.dot(At_v2, W1)[:,0].compute()

还请注意，运行 .persist 会将数组保留在内存中，因此如果内存是一个约束，您可能不想用 .persist 阻塞它。