TensorFlow后端Keras多GPU LSTM模型添加Masking层后运行报错

你遇到的这个问题，是旧版multi_gpu_model工具和Masking层结合时的典型兼容性bug——多GPU复制模型时，Masking层生成的掩码相关操作在GPU上找不到对应的内核实现，导致设备分配失败。下面给你几个实用的解决方案：

方案1：改用TensorFlow官方推荐的分布式策略（MirroredStrategy） #

multi_gpu_model已经被TensorFlow官方弃用，现在更推荐使用原生的MirroredStrategy，它对各类层的兼容性更好，能自动处理多GPU下的设备调度问题。

修改后的完整代码示例：

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Masking, LSTM, Dropout, TimeDistributed, Dense

# 自动检测并使用所有可用GPU
strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
    lstm_model = Sequential()
    lstm_model.add(Masking(mask_value=-5, input_shape=(Time_Steps, Num_Features)))
    lstm_model.add(LSTM(units=100, return_sequences=True))
    lstm_model.add(Dropout(0.2))
    lstm_model.add(LSTM(units=100, return_sequences=True))
    lstm_model.add(Dropout(0.2))
    lstm_model.add(TimeDistributed(Dense(25, activation='relu')))
    lstm_model.add(TimeDistributed(Dense(1, activation='sigmoid')))
    lstm_model.summary()
    
    lstm_model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])

# 直接用fit方法，无需再用multi_gpu_model包装
lstm_model.fit(
    x=training_generator,
    steps_per_epoch=Bbatches_afterGpuDiv,
    epochs=Epochs,
    verbose=1,
    validation_data=validation_generator,
    validation_steps=2,
    use_multiprocessing=True,
    workers=8,
    max_queue_size=8
)

方案2：手动在数据生成阶段处理掩码（替代Masking层） #

如果暂时不想切换分布式策略，可以在数据生成器里提前处理无效时间步，绕过Masking层的GPU兼容性问题：

1. 修改quickdrawSequence的__data_generation方法：

def __data_generation(self, curnt_batchFile):
    curnt_batchFile_x = curnt_batchFile
    curnt_batchFile_y = curnt_batchFile.replace("_x.npy","_y.npy")
    x_val = np.load(curnt_batchFile_x)
    y_val = np.load(curnt_batchFile_y)
    
    # 生成掩码：有效时间步标记为1，无效（值为-5）标记为0
    mask = np.all(x_val == -5, axis=-1).astype(np.float32)
    mask = 1 - mask  # 反转后，有效步为1，无效为0
    
    # 将无效时间步的特征值替换为0，避免干扰LSTM计算
    x_val = np.where(x_val == -5, 0, x_val)
    
    return [x_val, mask], y_val

2. 改用函数式API构建模型，手动传入掩码：

with tf.device('/cpu:0'):
    inputs = tf.keras.Input(shape=(Time_Steps, Num_Features))
    mask_input = tf.keras.Input(shape=(Time_Steps,))
    
    x = LSTM(units=100, return_sequences=True)(inputs, mask=mask_input)
    x = Dropout(0.2)(x)
    x = LSTM(units=100, return_sequences=True)(x)
    x = Dropout(0.2)(x)
    x = TimeDistributed(Dense(25, activation='relu'))(x)
    outputs = TimeDistributed(Dense(1, activation='sigmoid'))(x)
    
    lstm_model = tf.keras.Model(inputs=[inputs, mask_input], outputs=outputs)
    lstm_model.summary()

parallel_model = multi_gpu_model(lstm_model, gpus=4)
parallel_model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])

方案3：升级TensorFlow版本 #

如果你的TensorFlow还是1.x版本，建议升级到2.x的稳定版本——新版TensorFlow对GPU内核的支持更完善，很多旧版的兼容性问题已经被修复。

问题根源说明 #

当使用multi_gpu_model复制模型到多个GPU时，Masking层生成的掩码张量会触发transpose等操作，而旧版TensorFlow没有为这些操作提供对应的GPU内核实现，导致无法在GPU上执行，最终抛出设备分配错误。MirroredStrategy作为原生分布式策略，会自动处理这类跨设备的张量操作兼容性问题，是目前更可靠的多GPU训练方案。

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