TensorFlow中channels_first格式下LSTM层形状秩错误的解决方法

问题描述 #

我尝试使用TensorFlow Keras构建基于多模态3D图像的模型，输入形状为(2,128,128,64)（channels_first格式），相关代码如下：

from tensorflow import keras
from tensorflow.keras.layers import Conv3D, MaxPool3D, GlobalAveragePooling3D
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D,RepeatVector,Bidirectional,LSTM,TimeDistributed,Dropout
from tensorflow.keras.models import Model
from tensorflow.keras.losses import CategoricalCrossentropy

tf.keras.backend.set_image_data_format('channels_first')   

inputs = keras.Input((2, 128, 128, 64))

x = Conv3D(filters=16, kernel_size=3, activation='relu' , name='conv1')(inputs) #32
x = Conv3D(filters=32, kernel_size=3, activation="relu",name='conv2')(x) #64
x = MaxPool3D(pool_size=2)(x)
x = GlobalAveragePooling3D()(x)
x = Dense(1024,activation='relu')(x)
x = Dropout(0.3)(x)

dense2 = Dense(2, activation='softmax')(x) 

x = RepeatVector(2)(x)
LSTM1 = Bidirectional(LSTM(64,activation='relu',return_sequences=True,dropout=0.2,recurrent_dropout=0.2))(x)
LSTM1 = Bidirectional(LSTM(32, activation='relu',return_sequences=True,dropout=0.2,recurrent_dropout=0.2))(LSTM1)
dense4 = TimeDistributed(Dense(1024, activation='relu'))(LSTM1)
dense1 = TimeDistributed(Dense(3, activation='softmax'))(dense4) 
model = Model(inputs, outputs=[dense1,dense2])

运行时LSTM层出现如下ValueError错误：

ValueError: Exception encountered when calling layer "lstm_cell_1" (type LSTMCell).

Shape must be at least rank 3 but is rank 2 for '{{node bidirectional/forward_lstm/lstm_cell_1/BiasAdd}} = BiasAdd[T=DT_FLOAT, data_format="NCHW"](bidirectional/forward_lstm/lstm_cell_1/MatMul, bidirectional/forward_lstm/lstm_cell_1/split_1)' with input shapes: [?,64], [64].

Call arguments received:
  • inputs=tf.Tensor(shape=(None, 1024), dtype=float32)
  • states=('tf.Tensor(shape=(None, 64), dtype=float32)', 'tf.Tensor(shape=(None, 64), dtype=float32)')
  • training=False

我看到相关帖子建议使用channels_last格式，但想了解是否能在保持channels_first格式的前提下训练模型并解决该错误？

错误原因 #

全局设置tf.keras.backend.set_image_data_format('channels_first')会将所有Keras层的默认数据格式改为NCHW，包括LSTM内部的底层操作（比如BiasAdd）。而LSTM层的输入逻辑是基于(batch_size, timesteps, features)的3维格式，其内部计算会生成2维张量，与NCHW格式要求的至少3维输入冲突，从而触发维度错误。

解决方案（保持channels_first） #

不需要全局设置数据格式，而是在卷积相关层显式指定data_format='channels_first'，这样既让3D卷积、池化层使用channels_first处理输入，又不会影响LSTM等序列层的正常运行。

修改后的代码如下：

from tensorflow import keras
from tensorflow.keras.layers import Conv3D, MaxPool3D, GlobalAveragePooling3D
from tensorflow.keras.layers import Dense, RepeatVector, Bidirectional, LSTM, TimeDistributed, Dropout
from tensorflow.keras.models import Model
from tensorflow.keras.losses import CategoricalCrossentropy

# 去掉全局数据格式设置，改为在卷积层显式指定
inputs = keras.Input((2, 128, 128, 64))

# 每个卷积、池化、全局池化层添加data_format='channels_first'
x = Conv3D(filters=16, kernel_size=3, activation='relu', name='conv1', data_format='channels_first')(inputs)
x = Conv3D(filters=32, kernel_size=3, activation="relu", name='conv2', data_format='channels_first')(x)
x = MaxPool3D(pool_size=2, data_format='channels_first')(x)
x = GlobalAveragePooling3D(data_format='channels_first')(x)

x = Dense(1024, activation='relu')(x)
x = Dropout(0.3)(x)

dense2 = Dense(2, activation='softmax')(x) 

x = RepeatVector(2)(x)
LSTM1 = Bidirectional(LSTM(64, activation='relu', return_sequences=True, dropout=0.2, recurrent_dropout=0.2))(x)
LSTM1 = Bidirectional(LSTM(32, activation='relu', return_sequences=True, dropout=0.2, recurrent_dropout=0.2))(LSTM1)
dense4 = TimeDistributed(Dense(1024, activation='relu'))(LSTM1)
dense1 = TimeDistributed(Dense(3, activation='softmax'))(dense4) 

model = Model(inputs, outputs=[dense1, dense2])

说明 #

• 卷积相关层（Conv3D、MaxPool3D、GlobalAveragePooling3D）通过data_format='channels_first'参数，明确按照(channels, dim1, dim2, dim3)的格式处理输入，符合你的需求。
• LSTM等序列层不受全局数据格式影响，按照默认的(batch_size, timesteps, features)格式处理输入，避免了维度冲突。

内容的提问来源于stack exchange，提问作者Dushi Fdz