TensorFlow添加自定义层后训练极慢 算子跑在CPU而非GPU原因求解

• 你自定义层call方法中使用了三层Python原生for循环遍历batch、时间步、输入维度，TensorFlow Eager模式下Python侧控制流不会被自动调度到GPU执行，所有关联算子都会被默认放在CPU侧运行，这是张量不在GPU上的核心原因。
• 硬编码batch_size作为循环边界，导致TensorFlow静态图编译（@tf.function）无法适配动态输入，触发执行逻辑回退到CPU。
• 你在循环中插入的tf.print调试代码会强制将张量拷贝到CPU输出内容，进一步拖拽张量落到CPU运行。

1. 完全替换Python原生for循环为TensorFlow内置向量化操作：你当前的逻辑是实现多维度张量的笛卡尔外积计算，可以用tf.broadcast_to+tf.expand_dims组合替代循环内reshape+matmul的逻辑，用维度广播直接处理batch、时间步维度，完全消除Python侧控制流。
2. 移除硬编码的batch_size参数，改为在call方法中通过tf.shape(input_)动态获取输入的batch维度，适配静态图编译要求。
3. 给call方法添加@tf.function(jit_compile=True)装饰器，强制将前向逻辑编译为静态计算图，自动调度算子到GPU执行，XLA编译还能额外提升计算速度。
4. 删除call方法中的tf.print调试代码，避免张量被强制拷贝到CPU。

import tensorflow as tf
from tensorflow import keras

class RuleLayer(keras.layers.Layer):
    def __init__(self, time_steps, n_input, n_memb, **kwargs):
        super(RuleLayer, self).__init__(**kwargs)
        self.ts = time_steps
        self.n = n_input
        self.m = n_memb

    def build(self, batch_input_shape):
        super(RuleLayer, self).build(batch_input_shape)
        
    @tf.function(jit_compile=True)
    def call(self, input_):
        # 动态获取batch大小
        batch_size = tf.shape(input_)[0]
        # 初始化第一个维度的张量
        cp = input_[..., 0:1]
        # 广播外积计算替代循环
        for d in range(1, self.n):
            # 扩展当前cp的最后维度，广播输入的d维度做外积
            cp = tf.expand_dims(cp, axis=-1) * tf.expand_dims(input_[..., d], axis=-2)
        # 展平所有外积维度
        cp = tf.reshape(cp, (batch_size, self.ts, self.m ** self.n))
        return cp

    def compute_output_shape(self, batch_input_shape):
        return tf.TensorShape([batch_input_shape[0], self.ts, self.m ** self.n])
    
    def get_config(self):
        config = {
            'time_steps': self.ts,
            'n_input': self.n,
            'n_memb': self.m,
        }
        base_config = super(RuleLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

调用方式修改为去掉batch_size参数即可：

L2 = RuleLayer(420,4, 3, name='ruleLayer')(L1)

修改后所有算子都会被调度到GPU执行，单epoch耗时可降低到分钟级别。

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