5.4 TFX 流水线与自定义算子扩展

5.4 TFX 流水线与自定义算子扩展

本节摘要：一次训练成功靠人，一百次训练成功靠流水线。TFX 把"数据校验、特征工程、训练、评估、推送"固化成有向组件链，每步产物有版本、有元数据，训练不再是一次性脚本而是一条可审计的生产线。本节梳理 TFX 核心组件的分工与协作方式，给出用 Keras 模型接入流水线的最小配置思路；后半段回到框架底层，讲自定义层、自定义梯度与自定义算子三个扩展层级的边界。两者合起来回答一个问题：框架的默认件不够用时，你有哪些正规扩展点。

本节能力清单

阅读完本节，你应当能够：

1. 画出 TFX 主链各组件的输入输出关系，说出每步"防了什么事故"；
2. 描述 Keras 训练代码接入 TFX Trainer 的最小改造点；
3. 写一个自定义层并让它的梯度正确回传；
4. 区分"需要自定义梯度"与"需要自定义算子"的场景层级。

TFX 主链：每步防一种事故

TFX 组件链的设计逻辑是"把人工检查环节自动化、留痕化"。主链从上游到下游：

整条链的产物登记进元数据库：每次训练用了哪版数据、哪个特征定义、哪个模型通过评估，全部可查——机器学习工程的"审计需求"由此满足。

图 25 TFX 主链：数据到服务的流水线

Keras 接入 Trainer 的改造点

TFX Trainer 不关心你的建模代码长什么样，它要的是一份"run_fn 入口函数"：读入 transform 产物、建模、训练、按规范导出。改造点通常只有三处——数据从管道参数读、超参从配置读、导出走固定签名：

# run_fn 的骨架示意（TFX Trainer 的入口约定）
def run_fn(fn_args):
    import tensorflow as tf

    # 一、按 fn_args 给的路径读训练与评估数据（TFRecord）
    train_ds = tf.data.TFRecordDataset(fn_args.train_files).map(parse_and_transform).batch(64)
    eval_ds = tf.data.TFRecordDataset(fn_args.eval_files).map(parse_and_transform).batch(256)

    # 二、建模与训练（第 3 章的全套技术原样可用）
    model = build_model()
    model.fit(train_ds, validation_data=eval_ds,
              steps_per_epoch=fn_args.train_steps, verbose=0)

    # 三、按规范签名导出——Pusher 只认这个导出口
    model.export(fn_args.serving_model_dir)
    print("trainer fn done")
# 输出：trainer fn done
# 建模代码与流水线解耦：第 3 章的模型在第 5 章的流水线上原样运转
​

改造的心法是"参数化一切路径与步数"——流水线调度器会按数据版本与配置组合反复调用这份代码，硬编码的路径会让复现失效。

自定义层：表达缺失的结构

Keras 层库覆盖不了所有结构，自定义层是第一扩展层级。子类化 Layer，build 里建变量（首次调用时按输入形状创建）、call 里写前向。关键是让前向只用可微分的张量运算——梯度自动回传：

class AdaptiveDense(tf.keras.layers.Layer):
    """带可学习缩放门的自适应全连接层"""
    def __init__(self, units):
        super().__init__()
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(name="w",
                                 shape=(input_shape[-1], self.units),
                                 initializer="glorot_uniform")   # 首调用时创建
        self.gate = self.add_weight(name="gate", shape=(self.units,),
                                    initializer="ones")          # 门控初始全开

    def call(self, x):
        h = tf.matmul(x, self.w)
        return h * self.gate               # 逐单元缩放，梯度自动覆盖 gate

layer = AdaptiveDense(16)
out = layer(tf.random.normal([4, 8]))
print(out.shape, len(layer.trainable_variables))
# 输出：(4, 16) 2
# w 与 gate 都在可训练名单里，标准训练流程无缝接入
​

自定义梯度与自定义算子：两个更深的层级

第二层级是自定义梯度：前向用现成运算组合，但某一步的求导规则想手工指定（如直通估计器、梯度裁剪的定制策略）。tf.custom_gradient 装饰一个"前向加反向"函数对：

@tf.custom_gradient
def signed_sqrt(x):
    """符号平方根：前向 sign 乘 sqrt，反向手工定义"""
    def grad(dy):
        return dy / (tf.sqrt(tf.abs(x) + 1e-8) * 2.0)   # 自定义导数
    return tf.sign(x) * tf.sqrt(tf.abs(x)), grad

y = signed_sqrt(tf.constant([4.0, -9.0]))
print(y.numpy())
# 输出：[ 2. -3.]
with tf.GradientTape() as tape:
    t = tf.constant([4.0])
    tape.watch(t)
    z = signed_sqrt(t)
print(float(tape.gradient(z, t)[0]))
# 输出：0.25 —— 用的正是手工定义的导数规则
​

第三层级是自定义算子：前向本身没有现成张量运算可用（特殊硬件指令、高性能内核），要用 C++ 等底层语言写算子注册进框架。层级判断口诀：现成运算能组合出前向，自定义层就够；前向能写但默认求导不合适，上自定义梯度；前向本身没有张量运算可组合，才写自定义算子。三层的工程成本递增，绝大多数需求停在第一层。

⚠️ 常见坑：TFX 组件链里的 Transform 与线上推理的预处理逻辑不同源。Transform 的产出是一份"计算图形式的预处理"并随模型一起导出，这保证训练与线上逐字节一致——自己在线上重写一套预处理是最常见的训练线上偏差来源。

💡 关键直觉：扩展层级的选择标准是"改变哪一层"——改结构用自定义层，改求导用 custom_gradient，改内核才碰自定义算子。往下多走一层，维护成本翻一倍，能不上就不上。

本节要点回顾

• 组件链设计：每步防一种事故，产物登记元数据，交接只靠产物。
• Trainer 改造：run_fn 入口、参数化路径与步数、规范签名导出。
• 自定义层：build 建变量、call 写前向，可微分组合则梯度免费。
• 自定义梯度：前向反向成对提交，给特殊运算定制求导。
• 三级边界：层、梯度、算子成本递增，按"改哪层"对号入座。

第 5 章到站。第 6 章交出工具箱。