如何在CPU上使用经CudnnLSTM训练的GPU模型？（TensorFlow1.6.0+Ubuntu16.04）

问题根源 #

tf.contrib.cudnn_rnn.CudnnLSTM是TensorFlow专门为GPU硬件优化的LSTM实现，它直接依赖NVIDIA的CuDNN加速库，完全不支持CPU环境运行。所以当你在CPU环境尝试导出或加载使用了该层的模型时，会因为找不到CuDNN相关的底层操作、变量处理逻辑而抛出错误。

解决方案 #

1. 替换CudnnLSTM为原生LSTM层（最稳妥方案）

在模型训练阶段，直接把CudnnLSTM替换成TensorFlow原生的LSTM实现，这样模型可以同时兼容GPU和CPU环境。

示例代码替换：

# 原来的CudnnLSTM代码
from tensorflow.contrib.cudnn_rnn.python.layers.cudnn_rnn import CudnnLSTM
cudnn_lstm = CudnnLSTM(num_layers=2, num_units=256)
outputs, states = cudnn_lstm(inputs)

# 替换为TF1.x原生多层LSTM
from tensorflow.contrib.rnn import LSTMCell, MultiRNNCell
# 构建多层LSTM单元
cells = MultiRNNCell([LSTMCell(256) for _ in range(2)])
# 用dynamic_rnn处理序列输入（支持变长序列）
outputs, states = tf.nn.dynamic_rnn(cells, inputs, dtype=tf.float32)

替换后重新训练模型，导出的模型就能在CPU环境正常推理了。

2. 迁移权重到CPU兼容LSTM（无需重新训练）

如果已经用CudnnLSTM训练好模型，不想重新训练，可以通过迁移权重的方式转换成CPU兼容版本：

• 加载GPU训练好的CudnnLSTM模型，提取其权重参数
• 构建结构完全一致的原生LSTM层
• 将CudnnLSTM的权重映射到原生LSTM的变量中（注意：两者权重格式有差异，CudnnLSTM的权重是[输入维度+隐藏维度, 4*隐藏维度]的合并格式，原生LSTM需要拆分为输入到隐藏、隐藏到隐藏两部分权重）
• 保存新的模型用于CPU环境

3. 导出时强制指定CPU设备（辅助技巧）

如果模型已经移除GPU专属操作，导出时可以强制指定CPU设备，确保生成的模型适配CPU环境：

import tensorflow as tf

# 加载训练好的模型 checkpoint
saver = tf.train.import_meta_graph('your_model.meta')
with tf.Session(config=tf.ConfigProto(device_count={'GPU': 0})) as sess:
    saver.restore(sess, 'your_model')
    # 导出为 SavedModel 格式
    builder = tf.saved_model.builder.SavedModelBuilder('./cpu_compatible_model')
    # 定义输入输出签名（替换成你的实际张量名称）
    input_tensor = tf.get_default_graph().get_tensor_by_name('input:0')
    output_tensor = tf.get_default_graph().get_tensor_by_name('output:0')
    predict_signature = tf.saved_model.signature_def_utils.predict_signature_def(
        inputs={'input': input_tensor},
        outputs={'output': output_tensor}
    )
    builder.add_meta_graph_and_variables(
        sess, [tf.saved_model.tag_constants.SERVING],
        signature_def_map={'predict': predict_signature}
    )
    builder.save()

注意：这种方法仅适用于模型中已无GPU专属操作的情况，若仍保留CudnnLSTM层则无效。

额外提示 #

• TensorFlow 1.6版本较老旧，条件允许的话建议升级到TF1.x最新版本（如1.15），对CPU/GPU兼容性的支持更好
• 如果必须用CudnnLSTM做GPU加速训练，那么训练和导出都要在GPU环境完成，后续CPU环境使用时必须替换为兼容LSTM并迁移权重

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