在循环内计算索引的 Python theano

Question

我安装了 Theano 库来提高计算速度，这样我就可以使用 GPU 的强大功能。

但是，在计算的内部循环中，会根据循环索引和几个数组的对应值计算一个新索引。

然后使用该计算出的索引来访问另一个数组的元素，该数组又用于另一个计算。

这是否太复杂以至于不能期望 Theano 有任何显着的加速？

那么让我换个角度来换个角度来表述我的问题。 这是 GPU 代码片段的示例。出于简洁的原因，省略了一些初始化。我可以在不显着增加计算时间的情况下将其转换为 Python/Theano 吗？

__global__ void SomeKernel(const cuComplex* __restrict__  data,
                                 float* __restrict__ voxels)

{

unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;

unsigned int idy = blockIdx.y * blockDim.y + threadIdx.y;

unsigned int pos = (idy * NX + idx);

unsigned int ind1 = pos * 3;
float x = voxels[ind1];
float y = voxels[ind1 + 1];
float z = voxels[ind1 + 2];

int m;

for (m = 0; m < M; ++m)
{
    unsigned int ind2 = 3 * m;

    float diff_x = x - some_pos[ind2];
    float diff_y = y - some_pos[ind2 + 1];
    float diff_z = z - some_pos[ind2 + 2];

    float distance = sqrtf(diff_x * diff_x
                         + diff_y * diff_y
                         + diff_z * diff_z);

    unsigned int dist = rintf(distance/some_factor);
    ind3 = m * another_factor + dist;

    cuComplex some_element = data[ind3];

    Main calculation starts, involving some_element.

Answer 1

不，我没有看到使用张量而不是 for 循环无法完成的事情。这应该意味着您可能会看到速度有所提高，但这实际上取决于应用程序。您也有 python+theano 的开销，尤其是来自类 C 代码的开销。

所以，而不是

for (m = 0; m < M; ++m)
{
    unsigned int ind2 = 3 * m;

    float diff_x = x - some_pos[ind2];
    float diff_y = y - some_pos[ind2 + 1];
    float diff_z = z - some_pos[ind2 + 2];

    float distance = sqrtf(diff_x * diff_x
                         + diff_y * diff_y
                         + diff_z * diff_z);

    unsigned int dist = rintf(distance/some_factor);
    ind3 = m * another_factor + dist;

    cuComplex some_element = data[ind3];
}

你可以做类似(在我头顶)的事情

diff_xyz = T.Tensor([x,y,z]).dimshuffle('x',0) - some_pos.reshape(-1,3)
distance = T.norm(diff_xyz)
dist = T.round(distance/some_factor)
data = data.reshape(another_factor,-1)
some_elements = data[:,dist]

看到了吗？没有更多的循环，因此 GPU 可以将其并行化。

However, inside the inner loop of the computation a new index is calculated, based on the loop index and corresponding values of a couple of arrays. (...) Is this too complicated to expect any significant speedups from Theano?

一般来说:这可以通过使用张量而不是循环来优化，只要循环索引与所需索引具有线性关系。然而，它需要一点创造力和按摩才能正确。

非线性关系也可以使用 Tensor.take() ，但我不敢保证它在 GPU 上的速度。我的直觉总是告诉我远离它，因为它可能太灵活而无法很好地优化。但是，当没有其他选择时，可以使用。