Как получить AdvancedSubtensor на GPU

Question

У меня есть некоторый подзадачник и по какой-то причине, Theano не может перенести его на GPU.

Некоторые примеры кода:

import numpy 
import theano 
import theano.printing 
import theano.compile.io 
import theano.compile.function_module 
import theano.tensor as T 
from theano.sandbox.cuda.basic_ops import as_cuda_ndarray_variable 

n_copies, n_cells = 5, 10 
P = T.constant(numpy.zeros((n_copies, n_cells), dtype="int32")) # (n_copies,n_cells) -> list of indices 

meminkey = T.fmatrix() # (batch,n_cells) 
meminkey = as_cuda_ndarray_variable(meminkey) 
i_t = T.ones((meminkey.shape[0],)) 
batches = T.arange(0, i_t.shape[0]).dimshuffle(0, 'x', 'x') # (batch,n_copies,n_cells) 
P_bc = P.dimshuffle('x', 0, 1) # (batch,n_copies,n_cells) 
meminkeyP = meminkey[batches, P_bc] # (batch,n_copies,n_cells) 
meminkeyP = as_cuda_ndarray_variable(meminkeyP) 

func = theano.function(inputs=[meminkey], outputs=[meminkeyP]) 
theano.printing.debugprint(func) 

Я добавил некоторые as_cuda_ndarray_variable сделать проблему более ясно, потому что на выходе, вы видите переводы GpuFromHost и HostFromGpu, что было бы избежать, если это может сделать AdvancedSubtensor на GPU , Вывод.

Using gpu device 0: GeForce GTX TITAN (CNMeM is disabled, CuDNN not available) 
GpuFromHost [id A] '' 5 
|AdvancedSubtensor [id B] '' 4 
    |HostFromGpu [id C] '' 1 
    | |<CudaNdarrayType(float32, matrix)> [id D] 
    |InplaceDimShuffle{0,x,x} [id E] '' 3 
    | |ARange{dtype='int64'} [id F] '' 2 
    | |TensorConstant{0} [id G] 
    | |Shape_i{0} [id H] '' 0 
    | | |<CudaNdarrayType(float32, matrix)> [id D] 
    | |TensorConstant{1} [id I] 
    |TensorConstant{[[[4 0 1 2..5 8 9 7]]]} [id J] 

Итак, почему Theano не может преобразовать это в GPU?

Кроме того, как я могу переписать код, который Theano выполнит для вычисления на графическом процессоре?

---

Связанные вопросы в Google Groups: here и here и here.

Answer 1

Хорошо, поэтому в сообщениях групп Google, которые я связал, довольно хорошо объяснено, почему это не работает. AdvancedSubtensor - это самая общая форма, которая работает со всеми сумасшедшими типами индексирующих вариантов. Тогда есть AdvancedSubtensor1, который работает только для определенного типа подмножества. Существует только версия GPU для AdvancedSubtensor1, а не для AdvancedSubtensor. Я не совсем понял причину, но об этом постоянно идет дискуссия.

AdvancedSubtensor1 может использоваться, когда имеется один список индексов. Однако в моем примере это не так. Общее обходное решение, которое вы видите, также в каком-то другом примере в столбце Google Groups, заключается в том, чтобы сначала сгладить массив и рассчитать индексы для сплющенного массива.

Большинство примеров работают с каким-то nonzero() или так, где вы также сглаживаете базовые аргументы в том же порядке, а затем получаете индексы для сплющенной версии.

Итак, вопрос в том, как применить это к моему коду?

На самом деле, есть более простое решение, в котором он будет использовать AdvancedSubtensor1, который я не понял сначала:

meminkeyP = meminkey[:, P] # (batch,n_copies,n_cells) 

Однако, прежде чем я понял, что я придумал общее решение, которое также работает для других случаев. Я преобразую свои индексы tuple (batches, P_bc) в индексы для сплющенной версии. Это делается с помощью этой функции:

def indices_in_flatten_array(ndim, shape, *args): 
    """ 
    We expect that all args can be broadcasted together. 
    So, if we have some array A with ndim&shape as given, 
    A[args] would give us a subtensor. 
    We return the indices so that A[args].flatten() 
    and A.flatten()[indices] are the same. 
    """ 
    assert ndim > 0 
    assert len(args) == ndim 
    indices_per_axis = [args[i] for i in range(ndim)] 
    for i in range(ndim): 
    for j in range(i + 1, ndim): 
     indices_per_axis[i] *= shape[j] 
    indices = indices_per_axis[0] 
    for i in range(1, ndim): 
    indices += indices_per_axis[i] 
    return indices 

Затем я использую это так:

meminkeyP = meminkey.flatten()[indices_in_flatten_array(meminkey.ndim, meminkey.shape, batches, P_bc)] 

Это похоже на работу.

И я получаю этот выход:

Using gpu device 0: GeForce GTX TITAN (CNMeM is disabled, CuDNN not available) 
GpuReshape{3} [id A] '' 11 
|GpuAdvancedSubtensor1 [id B] '' 10 
| |GpuReshape{1} [id C] '' 2 
| | |<CudaNdarrayType(float32, matrix)> [id D] 
| | |TensorConstant{(1,) of -1} [id E] 
| |Reshape{1} [id F] '' 9 
| |Elemwise{second,no_inplace} [id G] '' 8 
| | |TensorConstant{(1, 5, 10) of 0} [id H] 
| | |Elemwise{Mul}[(0, 0)] [id I] '' 7 
| | |InplaceDimShuffle{0,x,x} [id J] '' 6 
| | | |ARange{dtype='int64'} [id K] '' 4 
| | | |TensorConstant{0} [id L] 
| | | |Shape_i{0} [id M] '' 0 
| | | | |<CudaNdarrayType(float32, matrix)> [id D] 
| | | |TensorConstant{1} [id N] 
| | |InplaceDimShuffle{x,x,x} [id O] '' 5 
| |  |Shape_i{1} [id P] '' 1 
| |  |<CudaNdarrayType(float32, matrix)> [id D] 
| |TensorConstant{(1,) of -1} [id E] 
|MakeVector{dtype='int64'} [id Q] '' 3 
    |Shape_i{0} [id M] '' 0 
    |TensorConstant{5} [id R] 
    |TensorConstant{10} [id S] 

Малого тест случай:

def test_indices_in_flatten_array(): 
    n_copies, n_cells = 5, 4 
    n_complex_cells = n_cells/2 
    n_batch = 3 
    static_rng = numpy.random.RandomState(1234) 
    def make_permut(): 
    p = numpy.zeros((n_copies, n_cells), dtype="int32") 
    for i in range(n_copies): 
     p[i, :n_complex_cells] = static_rng.permutation(n_complex_cells) 
     # Same permutation for imaginary part. 
     p[i, n_complex_cells:] = p[i, :n_complex_cells] + n_complex_cells 
    return T.constant(p) 
    P = make_permut() # (n_copies,n_cells) -> list of indices 

    meminkey = T.as_tensor_variable(static_rng.rand(n_batch, n_cells).astype("float32")) 
    i_t = T.ones((meminkey.shape[0],)) # (batch,) 
    n_batch = i_t.shape[0] 
    batches = T.arange(0, n_batch).dimshuffle(0, 'x', 'x') # (batch,n_copies,n_cells) 
    P_bc = P.dimshuffle('x', 0, 1) # (batch,n_copies,n_cells) 
    meminkeyP1 = meminkey[batches, P_bc] # (batch,n_copies,n_cells) 
    meminkeyP2 = meminkey.flatten()[indices_in_flatten_array(meminkey.ndim, meminkey.shape, batches, P_bc)] 

    numpy.testing.assert_allclose(meminkeyP1.eval(), meminkeyP2.eval())