python - 如何用theano扫描矩阵的所有元素？

Question

长话短说:

什么是 theano.scan 等同于:

M = np.arange(9).reshape(3, 3)
for i in range(M.shape[0]):
    for j in range(M.shape[1]):
        M[i, j] += 5
M

可能(如果可行)不使用嵌套的 scan？

请注意，这个问题并不想具体说明如何将一个元素的操作应用于矩阵，而是更一般地说明如何使用 theano.scan 实现一个像上面这样的嵌套循环结构。

---

长版:

theano.scan(或在本例中等效于 theano.map)允许通过简单地向sequences 参数，类似

import theano
import theano.tensor as T
M = T.dmatrix('M')
def map_func(i, j, matrix):
    return matrix[i, j] + i * j
results, updates = theano.scan(map_func,
            sequences=[T.arange(M.shape[0]), T.arange(M.shape[1])],
            non_sequences=[M])
f = theano.function(inputs=[M], outputs=results)
f(np.arange(9).reshape(3, 3))
# 

这大致相当于以下形式的 python 循环:

M = np.arange(9).reshape(3, 3)
for i, j in zip(np.arange(M.shape[0]), np.arange(M.shape[1])):
    M[i, j] += 5
M

将 M 的对角线中的所有元素增加 5。

但是如果我想找到 theano.scan 等价于:

M = np.arange(9).reshape(3, 3)
for i in range(M.shape[0]):
    for j in range(M.shape[1]):
        M[i, j] += 5
M

可能没有嵌套扫描？

一种方法当然是展平矩阵，扫描展平的元素，然后 reshape 到原来的形状，用像

import theano
import theano.tensor as T
M = T.dmatrix('M')
def map_func(i, X):
    return X[i] + .5
M_flat = T.flatten(M)
results, updates = theano.map(map_func,
                              sequences=T.arange(M.shape[0] * M.shape[1]),
                              non_sequences=M_flat)
final_M = T.reshape(results, M.shape)
f = theano.function([M], final_M)
f([[1, 2], [3, 4]])

但是有没有更好的方法不涉及明确展平和 reshape 矩阵？

Answer 1

这是一个例子，说明如何使用嵌套的 theano.scan 调用来实现这种事情。 在此示例中，我们将数字 3.141 添加到矩阵的每个元素，以复杂的方式有效地模拟了 H + 3.141 的输出:

H = T.dmatrix('H')
def fn2(col, row, matrix):
    return matrix[row, col] + 3.141

def fn(row, matrix):
    res, updates = theano.scan(fn=fn2,
                               sequences=T.arange(matrix.shape[1]),
                               non_sequences=[row, matrix])
    return res

results, updates = theano.scan(fn=fn,
                               sequences=T.arange(H.shape[0]),
                               non_sequences=[H])
f = theano.function([H], results)
f([[0, 1], [2, 3]])
# array([[ 3.141,  4.141],
#        [ 5.141,  6.141]])

再举一个例子，让我们给矩阵的每个元素加上它的行索引和列索引的乘积:

H = T.dmatrix('H')
def fn2(col, row, matrix):
    return matrix[row, col] + row * col

def fn(row, matrix):
    res, updates = theano.scan(fn=fn2,
                               sequences=T.arange(matrix.shape[1]),
                               non_sequences=[row, matrix])
    return res

results, updates = theano.scan(fn=fn,
                               sequences=T.arange(H.shape[0]),
                               non_sequences=[H])
f = theano.function([H], results)
f(np.arange(9).reshape(3, 3))
# Out[2]:array([[  0.,   1.,   2.],
#               [  3.,   5.,   7.],
#               [  6.,   9.,  12.]])