python - Theano 的 Scan 函数复制 non_sequences 共享变量

Question

我正在尝试在 Theano 中为 CNN 网络实现自定义卷积层，为此我使用了扫描函数。这个想法是将新的卷积掩模应用于每个像素。

scan 函数编译正确，但由于某种原因我收到内存不足错误。调试(见下文)表明，non_sequences 变量会为循环的每个实例(每个像素)进行复制，这当然会耗尽我的 GPU 内存:

def convolve_location(index, input, bias):
    hsize = self.W.shape / 2
    t = T.switch(index[0]-hsize[0] < 0, 0, index[0]-hsize[0])
    l = T.switch(index[1]-hsize[1] < 0, 0, index[1]-hsize[1])
    b = T.switch(index[0]+hsize[0] >= input.shape[2], input.shape[2]-1, index[0]+hsize[0])
    r = T.switch(index[1]+hsize[1] >= input.shape[3], input.shape[3]-1, index[1]+hsize[1])

    r_image = (input[:, :, t:b, l:r] - input[:, :, index[0], index[1]][:, :, None, None]) ** 2
    r_delta = (bias[:, :, t:b, l:r] - bias[:, :, index[0], index[1]][:, :, None, None]) ** 2
    return T.sum(r_image*r_delta)

# # Define cost function over all pixels
self.inds = theano.shared(np.array([(i, j) for i in range(self.image_shape[2]) for j in range(self.image_shape[3])], dtype='int32'), borrow=True)
self.cost = T.sum(theano.scan(
    fn=convolve_location,
    outputs_info=None,
    sequences=[self.inds],
    non_sequences=[self.input, self.b],
    n_steps=np.prod(self.image_shape[-2:])
)[0])

这是调试器的输出:

MemoryError: alloc failed Apply node that caused the error: Alloc(TensorConstant{0.0}, TensorConstant{1025}, TensorConstant{2000}, TensorConstant{3}, TensorConstant{32}, TensorConstant{32}) Inputs types: [TensorType(float32, scalar), TensorType(int64, scalar), TensorType(int64, scalar), TensorType(int64, scalar), TensorType(int64, scalar), TensorType(int64, scalar)] Inputs shapes: [(), (), (), (), (), ()] Inputs strides: [(), (), (), (), (), ()] Inputs values: [array(0.0, dtype=float32), array(1025), array(2000), array(3), array(32), array(32)]

Debugprint of the apply node:  Alloc [@A] <TensorType(float32, 5D)> '' |TensorConstant{0.0} [@B] <TensorType(float32, scalar)>  |TensorConstant{1025} [@C] <TensorType(int64, scalar)>  |TensorConstant{2000} [@D] <TensorType(int64, scalar)>  |TensorConstant{3} [@E] <TensorType(int64, scalar)>  |TensorConstant{32} [@F] <TensorType(int64, scalar)>  |TensorConstant{32} [@F] <TensorType(int64, scalar)> Storage map footprint:
 - CudaNdarrayConstant{[[[[ 0.]]]]}, Shape: (1, 1, 1, 1), ElemSize: 4 Byte(s), TotalSize: 4 Byte(s)
 - Constant{18}, Shape: (1,), ElemSize: 8 Byte(s), TotalSize: 8.0 Byte(s)

 - TensorConstant{(1, 1) of 0}, Shape: (1, 1), ElemSize: 1 Byte(s), TotalSize: 1 Byte(s)
 - Constant{1024}, Shape: (1,), ElemSize: 8 Byte(s), TotalSize: 8.0 Byte(s)
 - Constant{-1}, Shape: (1,), ElemSize: 8 Byte(s), TotalSize: 8.0 Byte(s)
 - TensorConstant{32}, Shape: (1,), ElemSize: 8 Byte(s), TotalSize: 8.0 Byte(s)
 - Subtensor{:int64:}.0, Shape: (1024,), ElemSize: 4 Byte(s), TotalSize: 4096 Byte(s)
 - Constant{34}, Shape: (1,), ElemSize: 8 Byte(s), TotalSize: 8.0 Byte(s)
 - Constant{2}, Shape: (1,), ElemSize: 8 Byte(s), TotalSize: 8.0 Byte(s)
 - TensorConstant{[2000    3..  32   32]}, Shape: (4,), ElemSize: 8 Byte(s), TotalSize: 32 Byte(s)
 - Reshape{4}.0, Shape: (2000, 3, 32, 32), ElemSize: 4 Byte(s), TotalSize: 24576000 Byte(s)
 - TensorConstant{(1, 1, 1, 1) of 0}, Shape: (1, 1, 1, 1), ElemSize: 1 Byte(s), TotalSize: 1 Byte(s)
 - CudaNdarrayConstant{[[[[ 0.1]]]]}, Shape: (1, 1, 1, 1), ElemSize: 4 Byte(s), TotalSize: 4 Byte(s)
 - <TensorType(float32, matrix)>, Shape: (50000, 3072), ElemSize: 4 Byte(s), TotalSize: 614400000 Byte(s)

您可以看到输入显示为 1025x2000x3x32x32 张量，而原始张量的大小为 2000x3x32x32，1025 是 scan + 1 的迭代次数。

为什么每次迭代都会复制 non_sequences 变量而不是简单地重复使用，我该如何修复它？

编辑:

self.input 和 self.b 都是共享变量。 self.input在初始化时传递给类，而self.b在类内部定义如下:

self.b = theano.shared(np.zeros(image_shape, dtype=theano.config.floatX), borrow=True)

Answer 1

当第一次创建扫描时或在优化过程中的某个时刻，可能会创建具有该形状的符号Alloc。 但是，应该在优化过程的后期进行优化。

我们知道最近有一个与此相关的问题，现在应该在 Theano 的开发(“前沿”)版本中修复。事实上，我刚刚使用最近的开发版本尝试了您的代码片段(稍作编辑)，并且没有内存错误。此外，计算图中的任何位置都没有 5D 张量，这表明该错误确实已被修复。

最后，请注意，诸如卷积之类的操作并不是真正的循环操作，当使用 scan 表示时，可能会比使用现有的卷积操作之一慢得多。特别是，当循环的迭代不相互依赖时，scan 将无法有效地并行化。