如何在Keras中逐列分割张量以实现STFCN

Question

我想在Keras实施Spatiotemporal全卷积网络（STFCN）。我需要输入3D卷积输出的每个深度列，例如形状为(64, 16, 16)的张量，作为单独LSTM的输入。

为清楚起见，我的(64 x 16 x 16)张量为(channels, height, width)。我需要将张量（显式或隐式）拆分为16 * 16 = 256个形状为(64 x 1 x 1)的张量。

以下是STFCN文件中的图表，用于说明Spatio-Temporal模块。我上面描述的是“空间特征”之间的箭头。和＆＃39;时空模块＆＃39;。

如何在Keras中最好地实现这个想法？

Answer 1

您可以使用Keras tf.split图层

从Tensorflow使用Lambda

使用Lambda将形状(64,16,16)的张量分割为(64,1,1,256)，然后将所需的任何索引子集化。

import numpy as np
import tensorflow as tf
import keras.backend as K
from keras.models import  Model
from keras.layers import Input, Lambda

# input data
data = np.ones((3,64,16,16))

# define lambda function to split
def lambda_fun(x) : 
    x = K.expand_dims(x, 4)
    split1 = tf.split(x, 16, 2)
    x = K.concatenate(split1, 4)
    split2 = tf.split(x, 16, 3)
    x = K.concatenate(split2, 4)
    return x

## check thet splitting works fine
input = Input(shape= (64,16,16))
ll = Lambda(lambda_fun)(input)
model = Model(inputs=input, outputs=ll)
res = model.predict(data)
print(np.shape(res))    #(3, 64, 1, 1, 256)