Python theano с индексом, вычисленным внутри цикла

Я установил библиотеку Theano для увеличения скорости вычислений, чтобы я мог использовать мощность графического процессора.

Однако во внутреннем цикле вычисления вычисляется новый индекс на основе индекса цикла и соответствующих значений пары массивов.

Затем вычисляемый индекс используется для доступа к элементу другого массива, который, в свою очередь, используется для другого вычисления.

Неужели это слишком сложно ожидать каких-либо значительных ускорений от Theano?

Так позвольте мне перефразировать мой вопрос, наоборот. Ниже приведен пример фрагмента кода графического процессора. Некоторые из них исключены из соображений краткости. Можно ли перевести это на 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.

Ответ 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]

См? Больше нет циклов, поэтому графический процессор может параллелизировать это.

Однако во внутреннем цикле вычисления вычисляется новый индекс, основанный на индексе цикла и соответствующих значениях нескольких массивов. (...) Неужели это слишком сложно ожидать каких-либо значительных ускорений от Theano?

В целом: это можно оптимизировать, если индекс цикла имеет линейную зависимость с требуемым индексом, используя тензоры вместо петель. Однако для этого нужно немного творчества и массажа.

Нелинейные отношения также возможны с помощью Tensor.take(), но я не решаюсь ручаться за его скорость на графическом процессоре. Мое чувство кисти всегда велело мне держаться подальше от него, так как это, вероятно, слишком гибко, чтобы оптимизировать красиво. Тем не менее, его можно использовать, когда нет альтернатив.

Ответ 2

Графические процессоры не работают в произвольной памяти при работе со своей глобальной памятью. Я раньше не использовал Theano, но если ваши массивы все вписываются в локальную память, это будет быстрым, так как случайный доступ там не проблема. Если это глобальная память, хотя трудно предвидеть, какая производительность будет, но это будет далеко от нее полной мощности. В другом примечании, что-то об этом вычислении даже параллелизуемо? Графические процессоры только на самом деле преуспевают, когда многие из этих вещей происходят одновременно.