Theano切换行有效

Question

我有以下代码

output = T.switch(cond, a, b)

这里cond是一个(N,1)布尔张量,而a和b是(N, M)与数值张量M是相当大的.该条件以行方式运行.

我可以很容易地通过运行开关工作T.repeat()的cond,但是这是相当缓慢.有没有办法在我能够有效地使布尔变量cond决定是否a或b应该归还？

Answer 1

有没有办法可以有效地让 cond 中的 bool 决定是否应该返回 a 或 b ？

是的，你可以做

cond * a + (1-cond) * b

cond将被广播(N, M)成型。

这应该接近理论极限，即内存带宽：此操作需要读取N*M元素并写入N*M。

相反，我们读取2*N*M，但删除条件逻辑。

（我面前没有 Theano，所以我不确定它是否比 更快T.switch，但它应该尽可能好。另外，我会尝试转换为与 和相同cond的值）dtypeab

a如果您想就地更新，可以使用T.set_subtensor：

a = np.random.uniform(size=(N, M)).astype(np.float32)
b = np.random.uniform(size=(N, M)).astype(np.float32)

a = theano.shared(a)
b = theano.shared(b)

c = T.vector() # mostly 0, presumably (1-cond)

nz = T.nonzero(c)

s = T.set_subtensor(a[nz], b[nz])
fn = theano.function([c], [], updates=[(a, s)])

...

fn(1-cond)

它可能会也可能不会比第一种方法更快，具体取决于N和M其他因素。