argmax作为subtensor的结果

时间:2015-07-10 13:33:19

标签: numpy theano

我想使用argmax保持尺寸作为子传感器。我有:

m, argm = T.max_and_argmax(a, axis=axis, keepdims=True)

我希望在a中将这些值设置为零。即我需要使用T.set_subtensor。要使用它,我需要在a_sub指定a的子argm,但我不确定它是怎么样的。多个维度a_sub = a[argm]错误。

这应该成立:

a_sub == T.max(a, axis=axis)
a_sub.shape == T.max(a, axis=axis).shape

最后,我想做:

a = T.set_subtensor(a_sub, 0)

我目前的解决方案:

idx = T.arange(a.shape[axis]).dimshuffle(['x'] * axis + [0] + ['x'] * (a.ndim - axis - 1))
a = T.switch(T.eq(idx, argm), 0, a)

但是,a_sub = a[T.eq(idx, argm)]不起作用。

1 个答案:

答案 0 :(得分:2)

你需要使用Theano' advanced indexing features,不幸的是,numpy's advanced indexing不同。

这是一个做你想做的事的例子。

更新:现在使用参数化轴,但请注意axis不能是符号。

import numpy

import theano
import theano.tensor as tt

theano.config.compute_test_value = 'raise'

axis = 2

x = tt.tensor3()
x.tag.test_value = numpy.array([[[3, 2, 6], [5, 1, 4]], [[2, 1, 6], [6, 1, 5]]],
                               dtype=theano.config.floatX)

# Identify the largest value in each row
x_argmax = tt.argmax(x, axis=axis, keepdims=True)

# Construct a row of indexes to the length of axis
indexes = tt.arange(x.shape[axis]).dimshuffle(
    *(['x' for dim1 in xrange(axis)] + [0] + ['x' for dim2 in xrange(x.ndim - axis - 1)]))

# Create a binary mask indicating where the maximum values appear
mask = tt.eq(indexes, x_argmax)

# Alter the original matrix only at the places where the maximum values appeared
x_prime = tt.set_subtensor(x[mask.nonzero()], 0)

print x_prime.tag.test_value