Theano - Среднее значение квадратов градиентов

Question

В anano, учитывая стоимость партии cost с формой (batch_size,), легко вычислить градиент средней стоимости, как в T.grad(T.mean(cost,axis=0),p), где p является параметром, используемым при вычислении cost. Это эффективно выполняется путем обратного распространения градиента через вычислительный граф. Теперь я хотел бы вычислить среднее значение квадратов градиентов над партией. Это можно сделать, используя следующий фрагмент кода:

import theano.tensor as T 

g_square = T.mean(theano.scan(lambda i:T.grad(cost[i],p)**2,sequences=T.arange(cost.shape[0]))[0],axis=0) 

Где для удобства p предполагается один тензор Theano, а не список тензоров. Вычисление может быть эффективно выполнено путем простого возврата градиента до последнего шага и возведения в квадрат компонентов последней операции (которая должна быть суммой по индексу партии). Возможно, я ошибаюсь в этом, но вычисление должно быть таким же простым и почти таким же быстрым, как простое обратное распространение. Тем не менее, theanano, похоже, не в состоянии оптимизировать вычисление, и он продолжает использовать цикл, делая вычисления чрезвычайно медленными.

Может кто-нибудь знать о решении сделать вычисления эффективными либо путем форсирования оптимизаций, выражая вычисления по-другому, либо даже проходя через процесс backpropagation?

Заранее спасибо.

Answer 1

Ваша функция g_square имеет сложность O (batch_size ** 2) вместо O (batch_size), как ожидалось. Это позволяет сделать его невероятно медленным для больших размеров партии.

Причина в том, что на каждой итерации передний и задний проходы вычисляются по всей партии, хотя требуется только cost[i] для одной точки данных. Я предполагаю, что входной график вычисления cost, x, является тензором с первым размером размера batch_size. Теано не имеет возможности автоматически срезать этот тензор вдоль этого измерения. Поэтому вычисление всегда выполняется по всей партии.

К сожалению, я не вижу лучшего решения, чем нарезка ввод и делает петлю снаружи Феано:

# x: input data batch 
batch_size = x.shape[0] 
g_square_fun = theano.function([p], T.grad(cost[0],p)**2) 

g_square_value = 0 
for i in batch_size: 
    g_square_value += g_square_fun(x[i:i+1]) 

Может быть, когда будущие версии Теано поставляются с более встраивать возможности для вычисления якобианы будет более элегантные решения.

Answer 2

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

# x: input data batch 
# assuming cost = network(x,p) 

from theano.gof.graph import clone_get_equiv 

def g_square(cost,p): 

    g = T.zeros_like(p) 

    def scan_fn(i, g, cost, p): 
     # clone the graph computing cost, but slice it's input 
     cloned = clone_get_equiv([],[cost], 
           copy_inputs_and_orphans=False, 
           memo={x: x[i:i+1]}) 
     cost_slice = cloned[cost].reshape([]) 
     return g+T.grad(cost_slice,p)**2 

    result,updates = theano.reduce(scan_fn, 
            outputs_info=g, 
            sequences=[T.arange(cost.size)], 
            non_sequences=[cost.flatten(),p]) 

    return result