TF2.0迁移：如何为tf.keras模型分层指定CPU/GPU设备部署？

嗨，这个场景我之前实操过，在TF2.0里要给不同层分配不同设备，子类化tf.keras.Model确实是最靠谱的方案——完全不用可惜Sequential，子类化的灵活性反而能让你精准控制每一层的运行设备～

核心思路 #

关键是在模型的call方法里，用tf.device()上下文管理器，把不同阶段的运算绑定到指定设备上。层的初始化不需要提前指定设备，TF会在第一次执行运算时，自动把层的变量创建在当前上下文指定的设备上，后续调用也会保持这个设备绑定。

完整代码示例 #

下面是针对你需求的实现：把Embedding层（含输入处理）放在CPU，后续LSTM、Dense等层放在GPU：

import tensorflow as tf

class SplitDeviceModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        # 先定义所有需要的层，不用提前指定设备
        self.embedding = tf.keras.layers.Embedding(input_dim=10000, output_dim=128)
        self.lstm = tf.keras.layers.LSTM(64)
        self.dense = tf.keras.layers.Dense(10, activation='softmax')

    def call(self, inputs):
        # 把输入和Embedding层的运算限定在CPU上
        with tf.device('/CPU:0'):
            embedded = self.embedding(inputs)
        
        # 把后续所有层的运算切换到GPU上
        with tf.device('/GPU:0'):
            lstm_out = self.lstm(embedded)
            final_output = self.dense(lstm_out)
        
        return final_output

# 初始化并验证模型
model = SplitDeviceModel()
# 告诉模型输入形状，让它自动构建权重
model.build(input_shape=(None, 100))
model.summary()

# 编译模型（根据你的任务调整优化器、损失函数）
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# 后续就可以正常用model.fit()训练了
# model.fit(train_data, train_labels, epochs=10, validation_split=0.2)

注意事项 #

• 设备可用性检查：先用tf.config.list_physical_devices('GPU')确认你的环境有可用GPU，不然指定/GPU:0会报错
• 跨设备张量传递：TF会自动处理CPU和GPU之间的张量复制，不用手动调用tf.transfer_to_device之类的方法，上下文切换时数据会自动迁移
• 多GPU场景：如果有多个GPU，可以指定具体设备，比如/GPU:1，只要和tf.config.list_physical_devices('GPU')返回的设备名对应就行
• 自定义层兼容：如果用自定义层，只要把自定义层的运算逻辑放到对应的tf.device()上下文里就行，逻辑完全一致

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