入门
安装与基础
TensorFlow 2.x 默认以 eager 模式运行(操作立即执行,如同普通 Python)。用 tf.config.list_physical_devices('GPU') 验证 CUDA GPU 是否可见。Linux/Windows 上安装 [and-cuda] 扩展以获得 GPU 支持,Apple Silicon 用 tensorflow-metal。TF 2.x 将 Keras 合并为 tf.keras。
# install TensorFlow (CPU)
pip install tensorflow
# with GPU support on Linux/Windows (CUDA-enabled)
pip install tensorflow[and-cuda]
import tensorflow as tf
# check version and devices
print(tf.__version__)
print("GPU:", tf.config.list_physical_devices('GPU'))
# eager execution is on by default in TF2
print(tf.executing_eagerly()) # True第一个张量
tf.constant 创建不可变张量;tf.zeros / tf.ones / tf.random.* 创建常见模式。每个张量有 shape、dtype 和 rank。.numpy() 可转回 NumPy。优先显式指定 dtype(float32 是 ML 默认值),避免静默类型提升导致性能下降。
import tensorflow as tf
# constants are immutable
x = tf.constant([[1, 2], [3, 4]])
y = tf.zeros((3, 3))
z = tf.random.normal((2, 3), mean=0.0, stddev=1.0)
print(x.shape, x.dtype) # (2, 2) int32
print(tf.reduce_sum(x)) # tf.Tensor(10, shape=(), dtype=int32)
# numpy interop is zero-copy when possible
print(x.numpy()) # array([[1, 2], [3, 4]])Eager 执行与梯度
Eager 执行让 TF 像普通 Python 一样可调试。GradientTape 记录被监听张量上的前向操作(tf.Variable 默认被监听),从而计算导数。这是自定义训练循环的基础。tape 默认一次性使用;用 persistent=True 可多次调用 gradient。
import tensorflow as tf
a = tf.constant(3.0)
b = tf.constant(4.0)
print(a * b + 1) # tf.Tensor(13.0, ...)
# GradientTape records ops for automatic differentiation
with tf.GradientTape() as tape:
tape.watch(a)
y = a ** 2 + b # y = a^2 + b, dy/da = 2a
grad = tape.gradient(y, a)
print(grad) # tf.Tensor(6.0, ...) (= 2*3)tf.function 与计算图
@tf.function 将 Python 函数转换为可移植、可优化的 TensorFlow 计算图——在 GPU 上快得多,且分布式训练必需。函数按输入签名追踪一次,因此避免内部使用 Python 副作用或数据相关的 Python 控制流。变量必须在函数外创建一次。
import tensorflow as tf
@tf.function # traces a graph; runs in graph mode after first call
def train_step(x, y):
with tf.GradientTape() as tape:
preds = x * w + b
loss = tf.reduce_mean((preds - y) ** 2)
grads = tape.gradient(loss, [w, b])
optimizer.apply_gradients(zip(grads, [w, b]))
return loss
w = tf.Variable(0.1)
b = tf.Variable(0.0)
optimizer = tf.keras.optimizers.SGD(0.01)
# first call builds the graph (slower); subsequent calls are fast
loss = train_step(tf.constant([1., 2.]), tf.constant([2., 4.]))混合精度
混合精度以 float16 计算但保留 float32 主副本,在 Tensor Core GPU(Volta+)上提供 2-3 倍吞吐和减半内存。强制输出层为 float32 以保证数值稳定性。需要计算能力 >= 7.0。
import tensorflow as tf
# enable mixed precision: compute in float16, store in float32
tf.keras.mixed_precision.set_global_policy('mixed_float16')
model = tf.keras.Sequential([
tf.keras.layers.Dense(64, activation='relu', dtype='float32'),
tf.keras.layers.Dense(10, activation='softmax', dtype='float32')
])
# loss scaling is automatic with mixed_float16 policy
print(tf.keras.mixed_precision.global_policy())张量基础
创建张量
tf.constant 用于固定数据;tf.convert_to_tensor 接受列表、NumPy 数组或标量。显式指定 dtype——默认推断可能为 Python 整数选 int32。tf.random.* 有 seed 参数用于复现;完全复现还需调用 tf.random.set_seed(seed)。
import tensorflow as tf
a = tf.constant([1, 2, 3])
b = tf.constant([[1.0, 2.0], [3.0, 4.0]], dtype=tf.float32)
zeros = tf.zeros((2, 3))
ones = tf.ones((3, 3))
eye = tf.eye(3)
full = tf.fill((2, 3), 7.0)
arange = tf.range(0, 10, 2) # [0, 2, 4, 6, 8]
linspace = tf.linspace(0.0, 1.0, 5)
rand = tf.random.uniform((2, 3))
randn = tf.random.normal((2, 3))
# constants are immutable: a[0] = 5 -> TypeError索引与切片
TF 支持 NumPy 风格切片,尽可能返回视图。tf.gather 沿一个轴选择;tf.gather_nd 处理多维索引数组。tf.boolean_mask 提取匹配布尔张量的元素。TF 张量不可变——用 tf.Variable 和 assign 进行原地更新。
import tensorflow as tf
w = tf.Variable(0.1, dtype=tf.float32)
W = tf.Variable(tf.random.normal((3, 4)))
b = tf.Variable(tf.zeros((4,)))
# read value
print(w.numpy()) # 0.1
# update value
w.assign(0.5)
w.assign_add(0.1) # w += 0.1
w.assign_sub(0.05)
# assign in place on a slice
W[0, 0].assign(99.0)
# watched by GradientTape automatically
with tf.GradientTape() as tape:
y = w ** 2
grad = tape.gradient(y, w) # 2w形状与重塑
t.shape 给出静态(编译时)形状;tf.shape(t) 返回动态(运行时)形状的张量——在 batch 大小可变的 tf.function 中使用。reshape 中 -1 自动推断一个维度。ndim>2 时 tf.transpose 需要 perm。用 expand_dims/squeeze 增删 size-1 轴以支持广播。
import tensorflow as tf
a = tf.constant([[1.0, 2.0], [3.0, 4.0]])
b = tf.constant([[5.0, 6.0], [7.0, 8.0]])
a + b; a - b; a * b; a / b # element-wise
c = tf.matmul(a, b) # matrix multiply
c = a @ b # equivalent
tf.reduce_sum(a) # sum of all
tf.reduce_sum(a, axis=0) # sum along axis 0
tf.reduce_mean(a, axis=1) # mean along axis 1
tf.reduce_max(a, axis=1)
tf.reduce_sum(a, axis=1, keepdims=True)
x = tf.range(12)
y = tf.reshape(x, (3, 4))
t = tf.transpose(y) # shape (4, 3)数学运算
标准运算符(+、*、@)已重载;tf.math.* 和 tf.linalg.* 覆盖其余。规约运算(reduce_sum/mean/max)接受 axis 和 keepdims=True 使结果仍可广播。用 tf.linalg 做矩阵运算(matmul、inv、svd、eigh、qr)。避免对张量使用 Python 循环——用 reduce/broadcast 运算向量化。
import tensorflow as tf
x = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
x[0] # first row
x[1, 2] # element (1,2) -> 6
x[:, 1] # second column
x[0:2] # first two rows
x[:, ::2] # every other column
# boolean mask
mask = x > 4
sel = tf.boolean_mask(x, mask) # tensor([5, 6, 7, 8, 9])
# gather by indices
idx = tf.constant([0, 2])
rows = tf.gather(x, idx) # rows 0 and 2
# where: select from two tensors
tf.where(x > 4, x, tf.zeros_like(x))广播与类型转换
广播遵循 NumPy 规则:维度从右对齐,size-1(或缺失)轴拉伸。tf.cast 显式改变 dtype——下转会截断。全程保持 float32 可避免 float64 带来的 2 倍内存和 GPU 减速。
import tensorflow as tf
# ragged: variable-length sequences
ragged = tf.ragged.constant([[1, 2, 3], [4], [5, 6]])
print(ragged.shape) # (3, None)
ragged.to_tensor() # pad to (3, 3)
# sparse: mostly-zero tensors
indices = [[0, 0], [1, 2], [2, 1]]
values = [1.0, 2.0, 3.0]
sparse = tf.sparse.SparseTensor(
indices=indices, values=values, dense_shape=[3, 3])
dense = tf.sparse.to_dense(sparse)
sp_sum = tf.sparse.reduce_sum(sparse)
sp_matmul = tf.sparse.sparse_dense_matmul(sparse, tf.eye(3))不规则张量与稀疏张量
RaggedTensor 处理变长 batch 无需手动填充——许多 tf.keras 层接受 ragged 输入。SparseTensor 存储 indices+values,用于 one-hot / TF-IDF 特征。StringTensor 保存字节串,配合 tf.strings.* 用于 tf.data 文本预处理。
import tensorflow as tf
# broadcasting follows NumPy rules
a = tf.constant([[1.0, 2.0, 3.0]]) # shape (1, 3)
b = tf.constant([[10.0], [20.0]]) # shape (2, 1)
c = a + b # shape (2, 3)
tf.abs(x); tf.sqrt(x); tf.exp(x); tf.log(x)
tf.square(x); tf.pow(x, 3)
tf.maximum(a, b); tf.minimum(a, b)
tf.clip_by_value(x, 0.0, 1.0)
tf.norm(x) # L2 norm
tf.norm(x, ord=1) # L1 norm
tf.cumsum(tf.range(5)) # [0, 1, 3, 6, 10]
tf.cumprod(tf.range(1, 5)) # [1, 2, 6, 24]变量与梯度
创建变量
tf.Variable 是 tf.constant 的可变对应物——通过 assign/assign_add/assign_sub 原地更新。可训练变量(trainable=True)被 GradientTape 收集。不可训练变量保存运行统计(BatchNorm)或步数计数器。
import tensorflow as tf
model = tf.keras.Sequential([
tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
tf.keras.layers.Dropout(0.5),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(10, activation='softmax'),
])
model.summary()
# build later with input_shape (or first call)
model = tf.keras.Sequential()
model.add(tf.keras.layers.Dense(128, activation='relu'))
model.add(tf.keras.layers.Dense(10))
model.build((None, 784)) # None = variable batch赋值与更新
用 assign/assign_add/assign_sub 做整张量更新;scatter_nd_update 修改特定索引。变量被 GradientTape 自动监听(常量需 tape.watch)。绝不在模型内用 v = ... 重新赋值——这会创建新 Python 变量并破坏梯度;始终用 v.assign(...)。
import tensorflow as tf
inputs = tf.keras.Input(shape=(784,), name='img')
x = tf.keras.layers.Dense(128, activation='relu')(inputs)
x = tf.keras.layers.Dropout(0.5)(x)
x = tf.keras.layers.Dense(64, activation='relu')(x)
outputs = tf.keras.layers.Dense(10, activation='softmax')(x)
model = tf.keras.Model(inputs=inputs, outputs=outputs, name='mlp')
# multi-input / multi-output
img_in = tf.keras.Input(shape=(224, 224, 3))
meta_in = tf.keras.Input(shape=(10,))
x = tf.keras.layers.Conv2D(32, 3, activation='relu')(img_in)
x = tf.keras.layers.GlobalAveragePooling2D()(x)
x = tf.keras.layers.concatenate([x, meta_in])
out = tf.keras.layers.Dense(1, activation='sigmoid')(x)
model = tf.keras.Model([img_in, meta_in], out)GradientTape 基础
GradientTape 记录被监听张量上的可微操作。tape.gradient 返回标量对目标的梯度。嵌套 tape 求高阶导数。用 tape.watch(x) 显式监听常量。梯度为 None 表示没有路径——常见的调试信号。
import tensorflow as tf
class MLP(tf.keras.Model):
def __init__(self, hidden=128, num_classes=10):
super().__init__()
self.dense1 = tf.keras.layers.Dense(hidden, activation='relu')
self.drop = tf.keras.layers.Dropout(0.5)
self.dense2 = tf.keras.layers.Dense(hidden, activation='relu')
self.out = tf.keras.layers.Dense(num_classes, activation='softmax')
def call(self, inputs, training=False):
x = self.dense1(inputs)
x = self.drop(x, training=training)
x = self.dense2(x)
return self.out(x)
model = MLP(hidden=64)
model.build((None, 784))
model.summary()持久 Tape 与雅可比
普通 tape 只能查询一次。persistent=True 保留 tape 以便多次调用 gradient()——记得 del 释放内存。tape.jacobian 计算完整雅可比矩阵(O(n*m) 内存);tape.batch_jacobian 高效处理每样本雅可比。
model.compile(
optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'],
loss_weights={'cls': 1.0, 'box': 0.5},
)
# explicit objects (more control)
model.compile(
optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
loss=tf.keras.losses.SparseCategoricalCrossentropy(),
metrics=[tf.keras.metrics.SparseCategoricalAccuracy(name='acc')],
run_eagerly=False, # debug mode: run without graph
)梯度操作
梯度裁剪抑制梯度爆炸:clip_by_global_norm 一起缩放所有梯度(首选),clip_by_value 逐元素硬截断。tf.stop_gradient 在一点切断计算图。调试时务必检查梯度的 NaN/Inf——它们会静默停滞训练。
model.summary()
print(model.layers)
print(model.inputs, model.outputs)
# plot architecture (needs pydot + graphviz)
tf.keras.utils.plot_model(
model, to_file='model.png', show_shapes=True,
show_layer_names=True, rankdir='TB',
)
print(f"{model.count_params():,} params")
config = model.get_config()
new_model = tf.keras.Model.from_config(config)非变量张量的梯度
常量需显式 tape.watch() 才能获得梯度。正常训练中模型的 Variables 自动被监听。tape.gradient 返回列表时,每项对应一个目标;None 表示没有可微路径——常见的静默训练停滞原因。
import tensorflow as tf
# preprocessing as part of the model (deployable)
model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(224, 224, 3)),
tf.keras.layers.Rescaling(1./255),
tf.keras.layers.RandomFlip('horizontal'),
tf.keras.layers.RandomRotation(0.1),
tf.keras.layers.Conv2D(32, 3, activation='relu'),
# ...
])
# the model now includes preprocessing, so inference needs raw images序贯模型
构建序贯模型
Sequential 是最简单的 Keras 模型:层的线性堆叠。在第一层传 input_shape(不含 batch 维)使 summary 在训练前可用。model.summary() 打印层形状和参数数。拓扑有分支时用函数式 API 或子类化。
import tensorflow as tf
dense = tf.keras.layers.Dense(
units=64, activation='relu', use_bias=True,
kernel_initializer='glorot_uniform',
bias_initializer='zeros',
kernel_regularizer=tf.keras.regularizers.l2(1e-4),
)
x = tf.random.normal((4, 128))
y = dense(x) # shape (4, 64)
print(dense.kernel.shape) # (128, 64)
print(dense.bias.shape) # (64,)
dense = tf.keras.layers.Dense(64, activation=tf.nn.gelu)编译与训练
compile() 配置优化器、损失和指标。标签为整数时用 sparse_categorical_crossentropy,one-hot 时用 categorical_crossentropy。fit() 返回 History 对象,其 .history 字典保存每轮损失和指标值。validation_split 切出验证集。
import tensorflow as tf
conv = tf.keras.layers.Conv2D(
filters=32, kernel_size=3, strides=(1, 1),
padding='same', activation='relu', use_bias=True,
)
x = tf.random.normal((4, 28, 28, 3)) # NHWC
out = conv(x)
print(out.shape) # (4, 28, 28, 32) with padding='same'
dw = tf.keras.layers.DepthwiseConv2D(3, padding='same', activation='relu')
up = tf.keras.layers.Conv2DTranspose(32, 3, strides=2, padding='same')逐层添加
Sequential.add() 逐个追加层。model.pop() 移除最后一层——迁移学习中常用。get_weights()/set_weights() 在 NumPy 数组列表间序列化权重。每层可通过 model.layers[i] 访问。
import tensorflow as tf
lstm = tf.keras.layers.LSTM(units=128, return_sequences=False)
x = tf.random.normal((4, 10, 64)) # (batch, time, features)
out = lstm(x) # (4, 128)
# stacked LSTM needs return_sequences=True
model = tf.keras.Sequential([
tf.keras.layers.LSTM(128, return_sequences=True, input_shape=(10, 64)),
tf.keras.layers.LSTM(64),
tf.keras.layers.Dense(10),
])
gru = tf.keras.layers.GRU(128)
bi = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(64))
out = bi(x) # shape (4, 128) — 64 * 2预测与评估
model.predict() 分批推理并返回 NumPy 数组。model.evaluate() 计算损失和指标。直接调用模型 model(x) 对单批更快,且可控制 training 标志:training=True 启用 Dropout 和 BatchNorm 更新。评估时始终设 training=False。
import tensorflow as tf
drop = tf.keras.layers.Dropout(0.5)
x = tf.random.normal((4, 10))
y = drop(x, training=True) # training=True activates dropout
tf.keras.layers.SpatialDropout2D(0.5) # drop entire feature maps
tf.keras.layers.GaussianDropout(0.5)
tf.keras.layers.AlphaDropout(0.5)
# weight regularization
tf.keras.layers.Dense(64,
kernel_regularizer=tf.keras.regularizers.l2(1e-4),
activity_regularizer=tf.keras.regularizers.l1(1e-5))
norm = tf.keras.layers.LayerNormalization()从 keras.Input 输入
添加 tf.keras.Input 作为第一层会立即构建模型,summary() 和 weights 即时可用。对于线性堆叠以外的拓扑——多输入、跳跃连接、共享层——切换到函数式 API。
import tensorflow as tf
bn = tf.keras.layers.BatchNormalization(
axis=-1, momentum=0.99, epsilon=1e-3,
)
model = tf.keras.Sequential([
tf.keras.layers.Conv2D(32, 3, padding='same', use_bias=False),
tf.keras.layers.BatchNormalization(),
tf.keras.layers.ReLU(),
])
print(len(bn.weights)) # 4
print(len(bn.trainable_weights)) # 2Custom Layer
Subclass Layer to build custom ops. Define weights in build() (called once with the input shape) rather than __init__ so the layer adapts to input size. Implement get_config() so the layer can be serialized and reloaded. add_weight registers the tensor with Keras' tracking (so it shows up in weights, is saved, and is moved to device). For stateless transforms, you can skip build() and just override call().
import tensorflow as tf
class Linear(tf.keras.layers.Layer):
def __init__(self, units=32, **kwargs):
super().__init__(**kwargs)
self.units = units
def build(self, input_shape):
self.w = self.add_weight(
shape=(input_shape[-1], self.units),
initializer='glorot_uniform', trainable=True, name='w')
self.b = self.add_weight(
shape=(self.units,), initializer='zeros', trainable=True, name='b')
def call(self, inputs):
return tf.matmul(inputs, self.w) + self.b
def get_config(self):
config = super().get_config()
config.update({'units': self.units})
return config
layer = Linear(64)
out = layer(tf.random.normal((4, 10)))函数式 API
多输入模型
函数式 API 将模型定义为层的 DAG——多输入/输出、残差连接或共享子模型必需。通过对另一个张量调用层 layer(x) 创建张量,然后用输入/输出张量实例化 Model。用按输入名称的字典训练。
import tensorflow as tf
bce = tf.keras.losses.BinaryCrossentropy(from_logits=False)
y_true = tf.constant([[1.0], [0.0], [1.0], [0.0]])
y_pred = tf.constant([[0.9], [0.1], [0.8], [0.2]])
loss = bce(y_true, y_pred)
# from_logits=True is more numerically stable
bce_logits = tf.keras.losses.BinaryCrossentropy(from_logits=True)
logits = tf.constant([[2.0], [-2.0], [1.5], [-1.5]])
loss = bce_logits(y_true, logits)
bce = tf.keras.losses.BinaryCrossentropy(
from_logits=True, pos_weight=tf.constant([5.0]))残差(跳跃)连接
残差连接将层输入加到输出:x_out = f(x) + x。这给梯度一条直接回传路径,防止梯度消失,使极深网络可训练(ResNet)。形状匹配时用 tf.keras.layers.Add();不同时用 1x1 Conv/Dense 投影。
import tensorflow as tf
# one-hot labels
cce = tf.keras.losses.CategoricalCrossentropy(from_logits=False)
y_true = tf.constant([[0, 1, 0], [1, 0, 0]], dtype=tf.float32)
y_pred = tf.constant([[0.1, 0.8, 0.1], [0.7, 0.2, 0.1]])
loss = cce(y_true, y_pred)
# integer labels -> use sparse version
scce = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
y_true_int = tf.constant([1, 0])
logits = tf.constant([[1.0, 3.0, 0.5], [2.5, 0.2, 0.1]])
loss = scce(y_true_int, logits)共享层与多输出
调用同一层实例两次即共享权重——两个输入经过相同 Embedding。多输出模型用损失字典编译(loss_weights 平衡任务)。fit 时传按输出名称的 y 字典。函数式 API 是表达共享层和多任务头的惯用方式。
import tensorflow as tf
mse = tf.keras.losses.MeanSquaredError()
y_true = tf.constant([3.0, -0.5, 2.0, 7.0])
y_pred = tf.constant([2.5, 0.0, 2.0, 8.0])
loss = mse(y_true, y_pred)
mae = tf.keras.losses.MeanAbsoluteError()
loss = mae(y_true, y_pred)
huber = tf.keras.losses.Huber(delta=1.0)
loss = huber(y_true, y_pred)
logcosh = tf.keras.losses.LogCosh()
loss = logcosh(y_true, y_pred)复用子模型
任何 tf.keras.Model 都可作为层在另一个模型内使用——组合块(编码器、解码器、骨干)的简洁方式。子模型的变量包含在父模型的 trainable_variables 中。冻结预训练子模型时在构建父模型前设 encoder.trainable = False。此模式支撑迁移学习。
import tensorflow as tf
def weighted_mse(y_true, y_pred):
weights = tf.cast(y_true != 0, tf.float32)
return tf.reduce_sum(weights * tf.square(y_true - y_pred)) / tf.reduce_sum(weights)
class FocalLoss(tf.keras.losses.Loss):
def __init__(self, alpha=0.25, gamma=2.0, **kwargs):
super().__init__(**kwargs)
self.alpha = alpha
self.gamma = gamma
def call(self, y_true, y_pred):
bce = tf.keras.losses.binary_crossentropy(y_true, y_pred)
p = y_pred
pt = p * y_true + (1 - p) * (1 - y_true)
return self.alpha * tf.pow(1 - pt, self.gamma) * bce
def get_config(self):
config = super().get_config()
config.update({'alpha': self.alpha, 'gamma': self.gamma})
return config
loss = FocalLoss(alpha=0.25, gamma=2.0)绘图与调试
model.summary() 验证形状和参数数。plot_model(需 pydot + graphviz)渲染图形 PNG——对记录有分支的函数式模型极有用。遍历 model.layers 可冻结或修改单个层。要检查中间激活,构建输出目标层的子 Model。
import tensorflow as tf
# why logits? numerical stability
logits = tf.constant([[100.0, 0.0, 0.0]])
# BAD: softmax first then CCE (can overflow)
probs = tf.nn.softmax(logits)
loss_bad = tf.keras.losses.CategoricalCrossentropy()(
tf.constant([[1.0, 0.0, 0.0]]), probs)
# GOOD: from_logits=True uses log-sum-exp internally
loss_good = tf.keras.losses.CategoricalCrossentropy(from_logits=True)(
tf.constant([[1.0, 0.0, 0.0]]), logits)
loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=y, logits=z)
loss = tf.nn.softmax_cross_entropy_with_logits(labels=y, logits=z)Multi-Output Losses
For multi-output models, pass a dict mapping output names to loss functions and a loss_weights dict to balance them. The total loss is the weighted sum. This is the standard pattern for multi-task learning — each head gets its own loss and metric. Pass training targets as a dict with the same keys at fit time. Tune loss_weights so no single task dominates the gradient.
import tensorflow as tf
model.compile(
optimizer='adam',
loss={
'priority': tf.keras.losses.BinaryCrossentropy(from_logits=True),
'department': tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
},
loss_weights={'priority': 1.0, 'department': 0.5},
metrics={'priority': ['accuracy'], 'department': ['accuracy']},
)
inputs = tf.keras.Input(shape=(64,))
x = tf.keras.layers.Dense(64, activation='relu')(inputs)
out_a = tf.keras.layers.Dense(1, name='priority')(x)
out_b = tf.keras.layers.Dense(10, name='department')(x)
model = tf.keras.Model(inputs, [out_a, out_b])
model.fit(x_train, {'priority': y_pri, 'department': y_dep}, epochs=5)模型子类化
基础子类化模型
子类化 tf.keras.Model 给予前向传播的完整 Python 控制——在 __init__ 定义层,在 call(self, inputs, training=False) 定义逻辑。training 标志切换 Dropout/BatchNorm。权重在首次调用时惰性创建,所以 summary/save 前需先调用一次。最灵活但运行时前无形状检查。
import tensorflow as tf
opt = tf.keras.optimizers.Adam(learning_rate=1e-3)
opt = tf.keras.optimizers.Adam(
learning_rate=1e-3, beta_1=0.9, beta_2=0.999,
epsilon=1e-7, amsgrad=False,
)
# schedule as learning_rate
lr_schedule = tf.keras.optimizers.schedules.CosineDecay(
initial_learning_rate=1e-3, decay_steps=10000)
opt = tf.keras.optimizers.Adam(learning_rate=lr_schedule)
# gradient clipping
opt = tf.keras.optimizers.Adam(learning_rate=1e-3, global_clipnorm=1.0)自定义层
子类化 tf.keras.layers.Layer 创建可复用的自定义层。在 build(input_shape) 中定义权重(形状已知时惰性调用)而非 __init__。add_weight 注册变量使其被 .trainable_weights 收集并自动保存。实现 get_config() 以支持序列化。
import tensorflow as tf
opt = tf.keras.optimizers.SGD(learning_rate=0.01)
opt = tf.keras.optimizers.SGD(learning_rate=0.01, momentum=0.9)
opt = tf.keras.optimizers.SGD(
learning_rate=0.01, momentum=0.9, nesterov=True)
# with weight decay (Keras 2.13+ / 3.x)
opt = tf.keras.optimizers.SGD(
learning_rate=0.01, momentum=0.9, weight_decay=1e-4)
# clip gradients by value
opt = tf.keras.optimizers.SGD(
learning_rate=0.01, momentum=0.9, clipvalue=0.5)带配置的自定义层
实现 get_config()(返回构造参数字典)后层可通过 model.to_json/from_json 序列化。始终调用 super().get_config() 并更新自己的字段。用 @tf.keras.utils.register_keras_serializable() 注册以在无显式类导入时重载。
import tensorflow as tf
opt = tf.keras.optimizers.AdamW(
learning_rate=3e-4, weight_decay=0.01,
beta_1=0.9, beta_2=0.999)
opt = tf.keras.optimizers.RMSprop(learning_rate=1e-3, rho=0.9)
opt = tf.keras.optimizers.Adagrad(learning_rate=1e-2)
opt = tf.keras.optimizers.Adadelta(learning_rate=1.0)
opt = tf.keras.optimizers.Nadam(learning_rate=1e-3)
opt = tf.keras.optimizers.Ftrl(learning_rate=0.1)多前向路径
子类化在残差栈等非线性拓扑中表现出色,同一块循环应用 N 次。call 内的 Python 控制流在 eager 模式下可行,但被 @tf.function 按字面追踪——数据相关控制流用 tf.cond/tf.while_loop。残差加需形状匹配。
import tensorflow as tf
opt = tf.keras.optimizers.Adam(1e-3)
vars = [w, b]
with tf.GradientTape() as tape:
preds = model(x)
loss = loss_fn(y, preds)
grads = tape.gradient(loss, vars)
grads, _ = tf.clip_by_global_norm(grads, 1.0)
opt.apply_gradients(zip(grads, vars))
for v, g in zip(vars, grads):
print(v.name, g is not None, g.numpy().norm() if g is not None else None)自定义训练步骤
重写 train_step 自定义 fit() 每批行为,同时保留 fit 的轮次循环、回调和分布式。self.compiled_loss/compiled_metrics 调用传给 compile 的内容。这是推荐的中间方案——添加对比损失、GAN 训练或自定义指标而无需重写整个循环。
import tensorflow as tf
sched = tf.keras.optimizers.schedules.PiecewiseConstantDecay(
boundaries=[500, 1500], values=[1e-3, 1e-4, 1e-5])
sched = tf.keras.optimizers.schedules.ExponentialDecay(
initial_learning_rate=1e-3, decay_steps=1000, decay_rate=0.96,
staircase=True)
sched = tf.keras.optimizers.schedules.CosineDecay(
initial_learning_rate=1e-3, decay_steps=10000, alpha=0.0)
opt = tf.keras.optimizers.Adam(learning_rate=sched)
print(opt.learning_rate(opt.iterations))Gradient Clipping & EMA
global_clipnorm (rescales all gradients together) is the most common choice and is what transformers use. clipvalue hard-clamps each element. For custom clipping, modify grads between tape.gradient and apply_gradients. Exponential Moving Average (EMA) of weights often gives a small accuracy bump at inference — apply it after each step and use the averaged weights for evaluation. TF Addons has a MovingAverage optimizer wrapper.
import tensorflow as tf
opt = tf.keras.optimizers.Adam(1e-3, global_clipnorm=1.0)
opt = tf.keras.optimizers.Adam(1e-3, clipvalue=0.5)
opt = tf.keras.optimizers.Adam(1e-3, clipnorm=1.0)
# manual clipping (more control)
grads = tape.gradient(loss, vars)
grads = [tf.clip_by_norm(g, 1.0) if g is not None else g for g in grads]
grads, _ = tf.clip_by_global_norm(grads, 1.0)
opt.apply_gradients(zip(grads, vars))
# exponential moving average of weights (improves generalization)
ema = tf.train.ExponentialMovingAverage(decay=0.999)层
Dense 与激活函数
Dense 是全连接主力:y = activation(x*W + b)。按任务选激活——隐藏层默认 ReLU,transformer 用 GELU,二分类用 sigmoid,多分类用 softmax。用 kernel_regularizer 加 L1/L2 权重衰减。初始化器很重要:tanh/sigmoid 用 glorot,ReLU 用 he。
import tensorflow as tf
model.compile(
optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'],
)
history = model.fit(
x_train, y_train, batch_size=32, epochs=10,
validation_data=(x_val, y_val),
validation_split=0.2, shuffle=True, verbose=2,
)
print(history.history['loss'])Conv2D 与池化
Conv2D 在空间维度上滑动学习滤波器。filters 是通道深度(每块翻倍:32->64->128)。padding='same' 零填充保持尺寸;'valid' 缩小。MaxPooling2D 下采样;GlobalAveragePooling2D 按平均折叠空间维度,比 Flatten 参数更少、过拟合更轻。
model.compile(
optimizer='adam', loss='sparse_categorical_crossentropy',
metrics=[
'accuracy',
tf.keras.metrics.SparseTopKCategoricalAccuracy(k=5, name='top5_acc'),
tf.keras.metrics.Precision(name='precision'),
tf.keras.metrics.Recall(name='recall'),
],
)
class F1Score(tf.keras.metrics.Metric):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.precision = tf.keras.metrics.Precision()
self.recall = tf.keras.metrics.Recall()
def update_state(self, y_true, y_pred, sample_weight=None):
self.precision.update_state(y_true, y_pred, sample_weight)
self.recall.update_state(y_true, y_pred, sample_weight)
def result(self):
p, r = self.precision.result(), self.recall.result()
return 2 * p * r / (p + r + 1e-8)
def reset_state(self):
self.precision.reset_state(); self.recall.reset_state()RNN / LSTM / GRU
LSTM/GRU 用门控记忆处理序列。return_sequences=True 返回完整序列(堆叠 RNN);False 返回最后时间步。Bidirectional 包装 RNN 双向读取——NLP 通常免费提升质量。GRU 参数更少。GPU 加速需保持默认激活以启用 CuDNN 内核。
model.fit(x_train, y_train, batch_size=32, epochs=10)
# with tf.data.Dataset
model.fit(train_dataset, epochs=10, validation_data=val_dataset)
# Dataset of unknown size: use steps_per_epoch
model.fit(
train_dataset, epochs=10, steps_per_epoch=500,
validation_data=val_dataset, validation_steps=50,
)
# initial epoch (resume training)
model.fit(train_dataset, initial_epoch=5, epochs=10)Dropout 与 BatchNorm
Dropout 训练时随机置零单元(rate = 丢弃概率)——推理时自动禁用。BatchNormalization 用 batch 统计将特征归一化为零均值/单位方差,再缩放/偏移——稳定训练、降低对初始化的敏感度。BatchNorm 放在激活前(偏置后);推理时用移动平均。
class_weight = {0: 1.0, 1: 5.0, 2: 1.0} # class 1 is rare
model.fit(x_train, y_train, class_weight=class_weight, epochs=10)
sample_weight = np.where(y_train == 1, 5.0, 1.0)
model.fit(x_train, y_train, sample_weight=sample_weight, epochs=10)
# with Dataset: yield (x, y, sample_weight) tuples
ds = tf.data.Dataset.from_tensor_slices((x, y, w))
model.fit(ds, epochs=10)
from sklearn.utils import class_weight
cw = class_weight.compute_class_weight('balanced', classes=np.unique(y), y=y)
class_weight = dict(enumerate(cw))Embedding 与序列
Embedding 是可学习查找表,将整数 token ID 映射为稠密向量。mask_zero=True 标记填充(id 0)使下游 RNN/注意力忽略它。Transformer 需在 embedding 上加位置编码(正弦或可学习),因为注意力本身是置换不变的。
preds = model.predict(x_test, batch_size=64)
print(preds.shape) # (N, num_classes)
pred = model.predict(x_test[:1])
loss, acc = model.evaluate(x_test, y_test, batch_size=64, verbose=2)
test_loss, test_acc, test_prec = model.evaluate(test_dataset)
for batch in test_dataset:
preds = model.predict_on_batch(batch[0])
preds = model(x_test, training=True) # callable form, test-time aug注意力与现代层
MultiHeadAttention 是 Transformer 核心——多头并行缩放点积注意力。自注意力时 query/key/value 传同一张量。Transformer 块用残差连接和 LayerNorm 包裹注意力 + 前馈网络。Pre-LN(子层前归一化)比 Post-LN 训练更稳定,无需 warmup。
import tensorflow as tf
class CustomModel(tf.keras.Model):
def train_step(self, data):
x, y = data
with tf.GradientTape() as tape:
y_pred = self(x, training=True)
loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)
grads = tape.gradient(loss, self.trainable_variables)
self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
self.compiled_metrics.update_state(y, y_pred)
return {m.name: m.result() for m in self.metrics}
def test_step(self, data):
x, y = data
y_pred = self(x, training=False)
self.compiled_loss(y, y_pred)
self.compiled_metrics.update_state(y, y_pred)
return {m.name: m.result() for m in self.metrics}
model = CustomModel(inputs, outputs)
model.compile(optimizer='adam', loss='...', metrics=['accuracy'])
model.fit(dataset, epochs=10)损失函数
分类损失
按标签格式匹配损失:整数标签用 SparseCategoricalCrossentropy,one-hot 用 CategoricalCrossentropy。用 from_logits=True 且不加 softmax 以保证数值稳定性。softmax 输出搭配 from_logits=True 会二次施加 softmax 并静默破坏训练。BinaryCrossentropy 用于 sigmoid 输出。
import tensorflow as tf
early = tf.keras.callbacks.EarlyStopping(
monitor='val_loss', min_delta=0.001, patience=5,
mode='auto', restore_best_weights=True, verbose=1,
)
model.fit(train_ds, validation_data=val_ds, epochs=100, callbacks=[early])回归损失
MSE(L2)是默认选择——梯度平滑但大误差被平方,异常值主导。MAE(L1)对异常值稳健但在 0 处不可微。Huber 兼具两者:零附近二次,远处线性。按噪声选择:高斯噪声用 MSE,重尾噪声用 MAE/Huber。
import tensorflow as tf
ckpt = tf.keras.callbacks.ModelCheckpoint(
filepath='checkpoints/epoch-{epoch:02d}-val-{val_loss:.3f}.keras',
monitor='val_loss', save_best_only=True, save_weights_only=False,
mode='min', save_freq='epoch', verbose=1,
)
ckpt = tf.keras.callbacks.ModelCheckpoint(
'weights.{epoch:02d}.h5', save_weights_only=True, save_freq='epoch',
)
model.fit(train_ds, validation_data=val_ds, epochs=50, callbacks=[ckpt])自定义损失函数
自定义损失是任何 (y_true, y_pred) -> 标量张量的可调用对象。可序列化损失需子类化 tf.keras.losses.Loss 并实现 call() 和 get_config()。只用 TF 运算(非 NumPy)以在 GradientTape 内可微。常见:不平衡类用 focal loss,分割用 Dice/IoU,嵌入用 triplet loss。
import tensorflow as tf
reduce = tf.keras.callbacks.ReduceLROnPlateau(
monitor='val_loss', factor=0.5, patience=3, min_lr=1e-7,
mode='min', min_delta=1e-4, cooldown=0, verbose=1,
)
model.fit(train_ds, validation_data=val_ds, epochs=100, callbacks=[reduce])损失权重与多任务
多输出模型用损失字典编译;loss_weights 平衡不同尺度的任务。fit 时传按输出名称的 y 字典。调 loss_weights 是多任务学习的关键超参——从逆量级权重开始,按哪个任务表现差来调整。
import tensorflow as tf
tb = tf.keras.callbacks.TensorBoard(
log_dir='./logs/fit', histogram_freq=1, write_graph=True,
write_images=False, update_freq='epoch',
profile_batch=2, embeddings_freq=0,
)
model.fit(train_ds, validation_data=val_ds, epochs=10, callbacks=[tb])
# launch: tensorboard --logdir=./logs正则化与附加损失
kernel_regularizer 惩罚权重(L2 = 权重衰减,L1 = 稀疏);activity_regularizer 惩罚层输出。call() 内的 layer.add_loss() 让层贡献任意损失(VAE 中的 KL),自动收集和添加——绝不要手动加。通过 model.losses 访问。
import tensorflow as tf
def schedule(epoch, lr):
if epoch < 10:
return lr
return lr * tf.math.exp(-0.1).numpy()
scheduler = tf.keras.callbacks.LearningRateScheduler(schedule, verbose=1)
class StepScheduler(tf.keras.callbacks.Callback):
def on_train_batch_begin(self, batch, logs=None):
step = self.model.optimizer.iterations.numpy()
new_lr = 1e-3 * (0.5 ** (step // 1000))
self.model.optimizer.learning_rate.assign(new_lr)
scheduler = tf.keras.callbacks.LearningRateScheduler(
lambda epoch: 1e-3 * 0.9 ** epoch)内置损失目录
Keras 内置丰富损失目录:hinge 变体用于最大间隔(SVM 风格),KLDivergence 用于分布匹配(VAE、蒸馏),MSLE 用于跨数量级的目标,CosineSimilarity 用于方向回归。NegativeLogLikelihood 配合概率输出层。多数有 from_logits 选项以保证稳定性。
import tensorflow as tf
class DebugCallback(tf.keras.callbacks.Callback):
def on_epoch_end(self, epoch, logs=None):
lr = self.model.optimizer.learning_rate
if hasattr(lr, 'numpy'):
lr = lr.numpy()
print(f"Epoch {epoch}: lr={lr}, logs={logs}")
def on_train_batch_end(self, batch, logs=None):
if batch % 100 == 0:
print(f" batch {batch}: loss={logs.get('loss'):.4f}")
def on_train_end(self, logs=None):
print("Training finished!")
model.fit(train_ds, epochs=5, callbacks=[DebugCallback()])优化器
带动量的 SGD
SGD + momentum(0.9)仍是大规模视觉训练的顶级优化器——泛化常优于 Adam。Nesterov 动量前瞻,通常有小幅提升。裁剪梯度(首选 global_clipnorm)防止爆炸。apply_gradients(zip(grads, vars)) 是自定义循环中的手动更新调用。
import tensorflow as tf
ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
def gen():
for i in range(len(x_train)):
yield x_train[i], y_train[i]
ds = tf.data.Dataset.from_generator(gen, output_signature=(
tf.TensorSpec(shape=(28, 28), dtype=tf.float32),
tf.TensorSpec(shape=(), dtype=tf.int32),
))
ds = ds.shuffle(10000).batch(32).prefetch(tf.data.AUTOTUNE)
for x, y in ds:
print(x.shape, y.shape)
breakAdam 及变体
Adam 从矩估计自适应每参数学习率——收敛快、调参少。AdamW 是 Adam 结合权重衰减的正确方式(普通 Adam + L2 不等价)。RMSprop 擅长 RNN。经验法则:transformer 用 AdamW,CNN 用 SGD+momentum,其他用 Adam 作为基线。
import tensorflow as tf
def augment(image, label):
image = tf.image.random_flip_left_right(image)
image = tf.image.random_brightness(image, 0.2)
return image, label
ds = tf.data.Dataset.from_tensor_slices((images, labels))
ds = ds.map(augment, num_parallel_calls=tf.data.AUTOTUNE)
ds = ds.map(lambda x, y: (tf.cast(x, tf.float32) / 255.0, y))
ds = ds.filter(lambda x, y: y < 10)
ds = ds.cache() # in memory
ds = ds.cache('/tmp/data.cache') # to file
ds = ds.cache().shuffle(1000).batch(32).prefetch(tf.data.AUTOTUNE)学习率调度
训练中衰减学习率至关重要。CosineDecay 平滑退火到 alpha*initial_lr——现代默认选择。Warmup(前 N 步线性升温)稳定大 batch 和 transformer 的早期训练;与余弦结合是标准 transformer 调度。直接将 Schedule 对象作为 learning_rate 传入。
import tensorflow as tf
ds = tf.data.Dataset.from_tensor_slices((x, y))
ds = ds.batch(32, drop_remainder=False)
ds = ds.shuffle(buffer_size=10000, seed=42, reshuffle_each_iteration=True)
ds = ds.repeat() # infinite
ds = ds.repeat(5) # 5 epochs
# recommended order: shuffle -> repeat -> batch
ds = ds.shuffle(10000).repeat().batch(32)自定义学习率调度
子类化 LearningRateSchedule 并实现 __call__(step) 构建任意调度——warmup+cosine 是事实上的 transformer 调度。用 tf.where 处理数据相关分支以在图中工作。实现 get_config() 以支持序列化。与 AdamW + 权重衰减组合是标准大模型配方。
import tensorflow as tf
ds = ds.prefetch(tf.data.AUTOTUNE)
ds = ds.map(parse_fn, num_parallel_calls=tf.data.AUTOTUNE)
# full recommended pipeline
ds = (tf.data.Dataset.from_tensor_slices((x, y))
.cache()
.shuffle(10000)
.map(augment, num_parallel_calls=tf.data.AUTOTUNE)
.batch(32)
.prefetch(tf.data.AUTOTUNE))梯度累积
梯度累积在 GPU 内存受限时模拟大 batch:对 N 个 micro-batch 求和梯度,除以 N,再应用一次。对大 batch LR 缩放 规则和在小 GPU 上拟合大模型至关重要。裁剪在累积后、优化器步骤前进行。部分 Keras 优化器现在原生支持 gradient_accumulation_steps。
import tensorflow as tf
ds = tf.data.TextLineDataset(['file1.txt', 'file2.txt'])
ds = tf.data.experimental.make_csv_dataset(
'data.csv', batch_size=32, label_name='label',
num_epochs=1, shuffle_buffer_size=1000,
)
ds = tf.keras.utils.image_dataset_from_directory(
'data/train', image_size=(224, 224), batch_size=32,
shuffle=True, seed=42, validation_split=0.2, subset='training',
)优化器状态与槽变量
自适应优化器在'槽变量'中保存每参数状态(Adam 的矩、动量缓冲),首次 apply_gradients 时创建并随模型保存——恢复训练时动量保持完整。通过 opt.learning_rate.assign(...) 原地调整学习率以保留状态;重建优化器会丢失状态。
import tensorflow as tf
def serialize(x, y):
feature = {
'x': tf.train.Feature(float_list=tf.train.FloatList(value=x.flatten())),
'y': tf.train.Feature(int64_list=tf.train.Int64List(value=[y])),
}
return tf.train.Example(features=tf.train.Features(feature=feature)).SerializeToString()
with tf.io.TFRecordWriter('data.tfrecord') as w:
for x, y in dataset:
w.write(serialize(x, y))
def parse(serialized):
feature = {'x': tf.io.FixedLenFeature([784], tf.float32),
'y': tf.io.FixedLenFeature([], tf.int64)}
parsed = tf.io.parse_single_example(serialized, feature)
return parsed['x'], parsed['y']
ds = tf.data.TFRecordDataset('data.tfrecord').map(parse).batch(32)回调
ModelCheckpoint 与 EarlyStopping
save_best_only=True 的 ModelCheckpoint 只保留监控指标上的最佳模型。EarlyStopping 在指标停止改善 `patience` 轮后停止训练;始终设 restore_best_weights=True 以最终获得最佳轮次的权重。两者结合:自动最佳模型选择。
import tensorflow as tf
model.save('saved_model') # SavedModel directory
loaded = tf.keras.models.load_model('saved_model')
loaded.evaluate(x_test, y_test)
# export for serving (explicit signature)
@tf.function(input_signature=[tf.TensorSpec([None, 784], tf.float32)])
def serve(x):
return {'outputs': model(x)}
tf.saved_model.save(model, 'serving_model', signatures=serve)ReduceLROnPlateau 与 CSV 记录
ReduceLROnPlateau 在 val_loss 停滞 `patience` 轮后减半学习率——无调度时的可靠默认选择。CSVLogger 将每轮指标写入 CSV。TerminateOnNaN 在损失变 NaN 时立即停止。三者组合是稳健的默认回调集。
import tensorflow as tf
model.save('model.keras') # Keras 3 default
loaded = tf.keras.models.load_model('model.keras')
# legacy HDF5 format
model.save('model.h5')
loaded = tf.keras.models.load_model('model.h5')
# .keras — Keras 3 default, supports custom layers via config
# .h5 — legacy, weights + architecture but limited custom layer support
# SavedModel — directory, best for deployment/servingTensorBoard
TensorBoard 记录损失/指标、权重直方图、计算图和性能分析数据用于浏览器实时可视化。histogram_freq=1 每轮追踪权重分布。profile_batch 分析一批的计算时间——性能调优极有用。用 `tensorboard --logdir <dir>` 启动,打开 localhost:6006。
import tensorflow as tf
model.save_weights('weights.keras')
model = build_model() # must match the saved architecture
model.load_weights('weights.keras')
model.save_weights('weights.h5')
model.load_weights('weights.h5')
# load weights from one model into another (partial)
model.load_weights('weights.keras', skip_mismatch=True, by_name=True)
model.save_weights('ckpt')
model.load_weights('ckpt')自定义回调
子类化 Callback 并重写钩子(on_epoch_begin/end、on_batch_begin/end)插入自定义逻辑。logs 是当前指标字典;self.model 可完全访问模型——设 self.model.stop_training = True 提前停止。回调可组合——传列表。
import tensorflow as tf
ckpt = tf.train.Checkpoint(model=model, optimizer=optimizer, step=optimizer.iterations)
manager = tf.train.CheckpointManager(
ckpt, directory='./ckpts', max_to_keep=3,
checkpoint_name='step-{step}',
)
manager.save()
ckpt.restore(manager.latest_checkpoint).expect_partial()
# in a custom loop: save every N steps
if step % 1000 == 0:
manager.save()学习率查找器
LR finder(Smith 2015)在一轮内指数增长 LR 同时记录平滑损失——最陡下降部分指示好的 LR。在子集上跑一轮,选损失下降最快处的 LR(比发散点低一个数量级)。绘制 self.lvs vs self.losses。
import tensorflow as tf
class MyLayer(tf.keras.layers.Layer):
...
custom_objects = {'MyLayer': MyLayer, 'focal_loss': FocalLoss}
loaded = tf.keras.models.load_model(
'model.keras', custom_objects=custom_objects,
)
# or register globally
@tf.keras.utils.register_keras_serializable(package='my_pkg')
class MyLayer(tf.keras.layers.Layer):
...回调组合
大多数生产运行使用相同回调集:EarlyStopping、ReduceLROnPlateau、ModelCheckpoint、TensorBoard、CSVLogger、TerminateOnNaN。它们干净组合并按列表顺序运行。微小一次性行为可用 LambdaCallback 传入 on_epoch_end/on_batch_end 可调用对象而无需写类。
import tensorflow as tf
# Keras 3 export API (recommended)
model.export('serving/1')
class ServingModule(tf.Module):
def __init__(self, model):
self.model = model
@tf.function(input_signature=[tf.TensorSpec([None, 784], tf.float32, name='inputs')])
def __call__(self, x):
return {'outputs': self.model(x)}
module = ServingModule(model)
tf.saved_model.save(module, 'serving/1',
signatures={'serving_default': module.__call__})
loaded = tf.saved_model.load('serving/1')
print(list(loaded.signatures.keys())) # ['serving_default']tf.data 数据管道
构建 Dataset
tf.data.Dataset 是标准输入抽象。from_tensor_slices 沿轴 0 将张量 zip 为 (x, y) 对。大数据集用 TFRecordDataset(二进制、分片)而非 from_tensor_slices(全加载入内存)。管道是惰性的——迭代或调用 fit 前不执行。
import tensorflow as tf
model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(28, 28, 1)),
tf.keras.layers.Conv2D(32, 3, padding='same', activation='relu'),
tf.keras.layers.BatchNormalization(),
tf.keras.layers.MaxPooling2D(), # 28 -> 14
tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu'),
tf.keras.layers.BatchNormalization(),
tf.keras.layers.MaxPooling2D(), # 14 -> 7
tf.keras.layers.Conv2D(128, 3, padding='same', activation='relu'),
tf.keras.layers.GlobalAveragePooling2D(), # (7,7,128) -> (128,)
tf.keras.layers.Dense(10, activation='softmax'),
])
model.summary()Map、Batch、Shuffle
标准管道顺序:shuffle、map、batch、prefetch。用大于数据集的缓冲区 shuffle 以充分混合。map 对每元素应用函数——用 num_parallel_calls=AUTOTUNE。map 后 batch 使每样本运算保持每样本。prefetch 使下一批准备与当前 GPU 计算重叠。
import tensorflow as tf
inputs = tf.keras.Input(shape=(32, 32, 3))
x = tf.keras.layers.Rescaling(1./255)(inputs)
for filters in [32, 64, 128]:
x = tf.keras.layers.Conv2D(filters, 3, padding='same', activation='relu')(x)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.Conv2D(filters, 3, padding='same', activation='relu')(x)
x = tf.keras.layers.MaxPooling2D()(x)
x = tf.keras.layers.Dropout(0.25)(x)
x = tf.keras.layers.GlobalAveragePooling2D()(x)
x = tf.keras.layers.Dense(256, activation='relu')(x)
x = tf.keras.layers.Dropout(0.5)(x)
outputs = tf.keras.layers.Dense(10, activation='softmax')(x)
model = tf.keras.Model(inputs, outputs)用 tf.image 增强
在 tf.data 内增强以提升性能——CPU 准备增强批,GPU 训练。用 tf.image.random_* 运算。保留单独的评估管道跳过 shuffle 和增强以使验证确定性。随机翻转对自然图像合适,但对数字/医学图像(方向重要)不适用。
import tensorflow as tf
import matplotlib.pyplot as plt
feature_model = tf.keras.Model(
inputs=model.inputs,
outputs=[layer.output for layer in model.layers if 'conv2d' in layer.name],
)
img = x_test[0:1]
features = feature_model.predict(img)
print([f.shape for f in features])
fig, axes = plt.subplots(4, 8, figsize=(12, 6))
for i, ax in enumerate(axes.flat):
if i < features[0].shape[-1]:
ax.imshow(features[0][0, :, :, i], cmap='viridis')
ax.axis('off')缓存与性能
cache() 存储管道到该点的输出——内存或磁盘。放在昂贵的确定性运算(解码、归一化)之后、随机运算(shuffle、增强)之前以免缓存随机性。大数据集(超内存)用磁盘缓存或 TFRecord。AUTOTUNE 动态选择并行度/预取大小。
import tensorflow as tf
backbone = tf.keras.applications.ResNet50(
include_top=False, weights='imagenet',
input_shape=(224, 224, 3), pooling='avg',
)
backbone.trainable = False
inputs = tf.keras.Input(shape=(224, 224, 3))
x = tf.keras.applications.resnet50.preprocess_input(inputs)
x = backbone(x, training=False)
x = tf.keras.layers.Dense(256, activation='relu')(x)
x = tf.keras.layers.Dropout(0.5)(x)
outputs = tf.keras.layers.Dense(10, activation='softmax')(x)
model = tf.keras.Model(inputs, outputs)TFRecord 读写
TFRecord 是 TF 的高效二进制格式——序列化 protobuf Example 顺序存储。推荐用于大数据集(图像、音频),因为比读取数千小文件管道效率高得多。用 TFRecordWriter 写,TFRecordDataset + parse_single_example 读。跨多文件分片以并行读取。
import tensorflow as tf
def residual_block(x, filters, stride=1):
shortcut = x
x = tf.keras.layers.Conv2D(filters, 3, strides=stride, padding='same', use_bias=False)(x)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.ReLU()(x)
x = tf.keras.layers.Conv2D(filters, 3, padding='same', use_bias=False)(x)
x = tf.keras.layers.BatchNormalization()(x)
if stride != 1 or shortcut.shape[-1] != filters:
shortcut = tf.keras.layers.Conv2D(filters, 1, strides=stride, use_bias=False)(shortcut)
shortcut = tf.keras.layers.BatchNormalization()(shortcut)
x = tf.keras.layers.Add()([x, shortcut])
return tf.keras.layers.ReLU()(x)
inputs = tf.keras.Input(shape=(32, 32, 3))
x = residual_block(inputs, 64)
x = residual_block(x, 128, stride=2)
model = tf.keras.Model(inputs, x)滑动窗口与时序
window(size, shift) 创建滑动窗口数据集(每个是嵌套数据集);flat_map + batch 展平。将时序转为监督 (输入, 目标) 对的标准模式。drop_remainder=True 仅保留完整窗口。同样适用于文本(预测下一 token)和任何序列建模任务。
import tensorflow as tf
# YOLO-style detection head: conv to (grid * grid * (5 + num_classes))
def detection_head(x, num_classes, num_anchors=3):
out = tf.keras.layers.Conv2D(
num_anchors * (5 + num_classes), 1, activation='linear')(x)
return out
# anchor-free alternative: predict (cx, cy, w, h, obj, class_logits)
def center_net_head(features, num_classes):
cls = tf.keras.layers.Conv2D(num_classes, 1, activation='sigmoid', name='cls')(features)
wh = tf.keras.layers.Conv2D(2, 1, name='wh')(features)
offset = tf.keras.layers.Conv2D(2, 1, name='offset')(features)
return {'cls': cls, 'wh': wh, 'offset': offset}训练与评估
用 Dataset 和类权重训练
class_weight 通过缩放损失重新平衡不平衡类别——'balanced' 设权重与频率成反比。不平衡的最简洁修复(无需合成数据、无需重采样)。sample_weight 更进一步,加权单个样本。两者都通过 tf.data yield (x, y, sample_weight) 元组集成。
import tensorflow as tf
lstm = tf.keras.layers.LSTM(128, return_sequences=False)
x = tf.random.normal((4, 10, 64))
out = lstm(x) # (4, 128)
gru = tf.keras.layers.GRU(128)
out = gru(x) # (4, 128)
lstm = tf.keras.layers.LSTM(128, return_sequences=True)
out = lstm(x) # (4, 10, 128)
bi = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(64))
out = bi(x) # (4, 128) — 64*2
lstm = tf.keras.layers.LSTM(64, stateful=True)evaluate 与 predict
evaluate 运行完整验证集并返回损失+指标。predict 返回 NumPy 输出。重写 predict_step 实现自定义推理如测试时增强(对增强副本取平均预测)。数据集长度未知时用 steps=。要批级控制,写自定义循环。
import tensorflow as tf
model = tf.keras.Sequential([
tf.keras.layers.Embedding(10000, 128, input_length=100, mask_zero=True),
tf.keras.layers.LSTM(128, return_sequences=True),
tf.keras.layers.Dropout(0.3),
tf.keras.layers.LSTM(64),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(10, activation='softmax'),
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')K 折交叉验证
k 折交叉验证比单次划分给出更稳健的泛化估计——小数据集必需。每折构建全新模型(不要复用权重)。每折内用 EarlyStopping + restore_best_weights。报告均值±标准差。大数据集通常单次留出足够。
import tensorflow as tf
model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(200,)),
tf.keras.layers.Embedding(10000, 128, mask_zero=True),
tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(64, return_sequences=True)),
tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(32)),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dropout(0.5),
tf.keras.layers.Dense(1, activation='sigmoid'),
])
model.compile(optimizer=tf.keras.optimizers.Adam(1e-3),
loss='binary_crossentropy', metrics=['accuracy'])
model.fit(train_ds, validation_data=val_ds, epochs=5)指标与混淆矩阵
传多个指标给 compile——AUC、Precision、Recall 是有状态的(跨批累积,每轮重置)。最终每类分析,在测试集上预测并喂给 sklearn 的 confusion_matrix 和 classification_report——给出每类精确率/召回率/F1,对诊断模型混淆哪些类极有价值。
import tensorflow as tf
encoder_input = tf.keras.Input(shape=(None,), name='enc_in')
enc_emb = tf.keras.layers.Embedding(vocab_src, 128)(encoder_input)
encoder = tf.keras.layers.LSTM(256, return_state=True)
enc_out, state_h, state_c = encoder(enc_emb)
decoder_input = tf.keras.Input(shape=(None,), name='dec_in')
dec_emb = tf.keras.layers.Embedding(vocab_tgt, 128)(decoder_input)
decoder_lstm = tf.keras.layers.LSTM(256, return_sequences=True, return_state=True)
dec_out, _, _ = decoder_lstm(dec_emb, initial_state=[state_h, state_c])
decoder_dense = tf.keras.layers.Dense(vocab_tgt, activation='softmax')
outputs = decoder_dense(dec_out)
model = tf.keras.Model([encoder_input, decoder_input], outputs)
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')混合精度训练
混合精度(float16 计算,float32 主权重)在现代 GPU 上给 2-3 倍加速和减半内存,但需要:(1) 设 mixed_float16 策略,(2) 强制最终输出为 float32,(3) 用 LossScaleOptimizer 包装优化器以防止 fp16 梯度下溢。'mixed_bfloat16' 无需损失缩放但旧 GPU 上较慢。
import tensorflow as tf
def make_dataset(series, window=20, batch=32):
ds = tf.data.Dataset.from_tensor_slices(series)
ds = ds.window(window + 1, shift=1, drop_remainder=True)
ds = ds.flat_map(lambda w: w.batch(window + 1))
ds = ds.map(lambda w: (w[:-1][..., tf.newaxis], w[-1:]))
return ds.shuffle(1000).batch(batch).prefetch(tf.data.AUTOTUNE)
model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(20, 1)),
tf.keras.layers.LSTM(32, return_sequences=False),
tf.keras.layers.Dense(1),
])
model.compile(optimizer='adam', loss='mse')
model.fit(make_dataset(train_series), epochs=10)分布式训练(MirroredStrategy)
MirroredStrategy 在单机多 GPU 上做同步数据并行训练:变量镜像,每步后梯度 all-reduce。在 strategy.scope() 内构建模型。全局 batch size 按 num_replicas_in_sync 缩放以保持每 GPU batch 相同。fit 自动处理分布式。
import tensorflow as tf
class BahdanauAttention(tf.keras.layers.Layer):
def __init__(self, units):
super().__init__()
self.W1 = tf.keras.layers.Dense(units)
self.W2 = tf.keras.layers.Dense(units)
self.V = tf.keras.layers.Dense(1)
def call(self, query, values):
q = tf.expand_dims(query, 1)
scores = self.V(tf.nn.tanh(self.W1(values) + self.W2(q)))
attn = tf.nn.softmax(scores, axis=1)
context = tf.reduce_sum(attn * values, axis=1)
return context, attn
attn = tf.keras.layers.AdditiveAttention()
context = attn([query, value])自定义训练循环
基础自定义循环
自定义循环给予完全控制:GradientTape 内前向、计算梯度、应用、手动更新指标。用 training=True 使 Dropout/BatchNorm 正确行为。指标有状态——每批 update_state,轮末读 result(),然后 reset_state()。用 @tf.function 包裹内循环可获得约 10 倍加速。
import tensorflow as tf
mha = tf.keras.layers.MultiHeadAttention(
num_heads=8, key_dim=64, dropout=0.1, use_bias=True)
x = tf.random.normal((4, 10, 512))
out, attn_weights = mha(query=x, value=x, key=x,
return_attention_scores=True)
print(out.shape) # (4, 10, 512)
print(attn_weights.shape) # (4, 8, 10, 10)
# causal mask (decoder self-attention)
mask = tf.linalg.band_part(tf.ones((10, 10)), -1, 0)
out = mha(query=x, value=x, key=x, attention_mask=mask)
# key padding mask (ignore pad positions)
padding = tf.cast(tokens == 0, tf.bool)[:, tf.newaxis, tf.newaxis, :]
out = mha(query=x, value=x, key=x, attention_mask=padding)tf.function 训练步骤
用 @tf.function 包裹 train_step 追踪为优化图——通常比 eager 快 5-10 倍。外层 Python 循环保持 eager(轮次迭代、日志)。别忘了 model.losses 以包含正则化损失。保持输入签名稳定以避免重新追踪;传张量而非 Python 整数/浮点数。
import tensorflow as tf
class TransformerBlock(tf.keras.layers.Layer):
def __init__(self, dim=512, heads=8, ff=2048, dropout=0.1, **kwargs):
super().__init__(**kwargs)
self.att = tf.keras.layers.MultiHeadAttention(num_heads=heads, key_dim=dim // heads)
self.ffn = tf.keras.Sequential([
tf.keras.layers.Dense(ff, activation='gelu'),
tf.keras.layers.Dense(dim),
])
self.norm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
self.norm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
self.drop1 = tf.keras.layers.Dropout(dropout)
self.drop2 = tf.keras.layers.Dropout(dropout)
def call(self, x, training=False, mask=None):
attn = self.att(x, x, x, attention_mask=mask)
x = self.norm1(x + self.drop1(attn, training=training))
ffn = self.ffn(x)
return self.norm2(x + self.drop2(ffn, training=training))自定义循环中的指标与验证
Mean 对标量损失跨批取平均;SparseCategoricalAccuracy 累积预测并计算整轮准确率。每轮开始 reset_state()(否则指标持续累积看起来人为平滑)。训练/验证保持各自追踪器以干净比较。这些簿记是完全控制的代价。
import tensorflow as tf
import numpy as np
class PositionalEncoding(tf.keras.layers.Layer):
def __init__(self, max_len=5000, dim=512, **kwargs):
super().__init__(**kwargs)
pos = np.arange(max_len)[:, np.newaxis]
i = np.arange(dim)[np.newaxis, :]
angle = pos / np.power(10000, (2 * (i // 2)) / np.float32(dim))
angle[:, 0::2] = np.sin(angle[:, 0::2])
angle[:, 1::2] = np.cos(angle[:, 1::2])
self.pos = tf.constant(angle[np.newaxis, ...], dtype=tf.float32)
def call(self, x):
return x + self.pos[:, :tf.shape(x)[1], :]
model = tf.keras.Sequential([
tf.keras.layers.Embedding(10000, 512),
PositionalEncoding(max_len=200, dim=512),
tf.keras.layers.Dropout(0.1),
])梯度累积循环
梯度累积在 GPU 内存受限时模拟大 batch:运行 N 次前向/反向,求和梯度,除以 N,应用一次。对大 batch LR 缩放规则和在小 GPU 上拟合大模型至关重要。裁剪在累积后、优化器步骤前进行。
import tensorflow as tf
class TransformerEncoder(tf.keras.layers.Layer):
def __init__(self, dim, heads, ff, num_layers, dropout=0.1, **kwargs):
super().__init__(**kwargs)
self.blocks = [TransformerBlock(dim, heads, ff, dropout)
for _ in range(num_layers)]
def call(self, x, training=False, mask=None):
for block in self.blocks:
x = block(x, training=training, mask=mask)
return x
inputs = tf.keras.Input(shape=(200,))
x = tf.keras.layers.Embedding(10000, 128, mask_zero=True)(inputs)
x = PositionalEncoding(200, 128)(x)
x = TransformerEncoder(128, 4, 512, 4)(x)
x = tf.keras.layers.GlobalAveragePooling1D()(x)
x = tf.keras.layers.Dropout(0.3)(x)
outputs = tf.keras.layers.Dense(1, activation='sigmoid')(x)
model = tf.keras.Model(inputs, outputs)分布式自定义循环
分布式自定义循环:在 strategy.scope() 内构建模型+优化器,分发数据集,用 strategy.run 包裹每副本步骤。关键是 loss 除以 num_replicas_in_sync 使 all-reduce 后梯度与单设备训练匹配。用 strategy.reduce(SUM) 聚合每副本损失以记录。
import tensorflow as tf
class ViTPatchEmbed(tf.keras.layers.Layer):
def __init__(self, patch_size=16, dim=768, **kwargs):
super().__init__(**kwargs)
self.proj = tf.keras.layers.Conv2D(dim, patch_size, strides=patch_size)
def call(self, x):
x = self.proj(x)
shape = tf.shape(x)
x = tf.reshape(x, [shape[0], -1, x.shape[-1]])
return x
inputs = tf.keras.Input(shape=(224, 224, 3))
x = ViTPatchEmbed(16, 768)(inputs)
x = x + tf.Variable(tf.random.normal((1, 196, 768)))
cls = tf.Variable(tf.random.normal((1, 1, 768)))
x = tf.keras.layers.Concatenate(axis=1)([cls, x])
for _ in range(6):
x = TransformerBlock(768, 12, 3072, 0.1)(x)
x = tf.keras.layers.Lambda(lambda t: t[:, 0])(x)
outputs = tf.keras.layers.Dense(10, activation='softmax')(x)
model = tf.keras.Model(inputs, outputs)GAN 训练循环
GAN 需自定义循环因为生成器和判别器交替更新,各有自己的优化器。每步两个 GradientTape——一个给 D(真 vs 假),一个给 G(欺骗 D)。非饱和生成器损失(最大化 log D(fake))提供更强早期梯度。Adam 用 beta_1=0.5(GAN 标准)。GAN 训练以不稳定著称。
import tensorflow as tf
class DecoderBlock(tf.keras.layers.Layer):
def __init__(self, dim, heads, ff, dropout=0.1, **kwargs):
super().__init__(**kwargs)
self.self_att = tf.keras.layers.MultiHeadAttention(heads, dim // heads)
self.cross_att = tf.keras.layers.MultiHeadAttention(heads, dim // heads)
self.ffn = tf.keras.Sequential([
tf.keras.layers.Dense(ff, activation='gelu'),
tf.keras.layers.Dense(dim),
])
self.norm1 = self.norm2 = self.norm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
self.drop = tf.keras.layers.Dropout(dropout)
def call(self, x, enc_out, training=False, causal_mask=None, padding_mask=None):
attn = self.self_att(x, x, x, attention_mask=causal_mask)
x = self.norm1(x + self.drop(attn, training=training))
cross = self.cross_att(query=x, value=enc_out, key=enc_out, attention_mask=padding_mask)
x = self.norm2(x + self.drop(cross, training=training))
x = self.norm3(x + self.drop(self.ffn(x), training=training))
return xCNN 图像分类
MNIST CNN 基线
标准小型 CNN 在 MNIST 上达约 99.4%。模式:Conv、BN、Conv、Pool、Dropout 块,滤波器翻倍(32->64),然后 GlobalAveragePooling + Dense 头。BatchNorm 稳定训练;Dropout 正则化。输入归一化到 [0,1] 并加通道轴。更难的数据集需更深(ResNet)和大量增强。
import tensorflow as tf
tb = tf.keras.callbacks.TensorBoard(
log_dir='logs/experiment_1', histogram_freq=1,
write_graph=True, profile_batch=2,
)
model.fit(train_ds, validation_data=val_ds, epochs=10, callbacks=[tb])
# launch tensorboard
# tensorboard --logdir=logs
import datetime
log_dir = 'logs/fit/' + datetime.datetime.now().strftime('%Y%m%d-%H%M%S')数据增强层
Keras 预处理层(RandomFlip、RandomRotation 等)在模型图内于 GPU 上运行——比 tf.data 增强更快且随模型保存。推理时自动停用。迁移学习时也应用骨干网络的 preprocess_input 以匹配预训练权重期望的归一化。
import tensorflow as tf
writer = tf.summary.create_file_writer('logs/custom')
with writer.as_default():
tf.summary.scalar('train/loss', loss, step=step)
tf.summary.scalar('train/lr', lr, step=step)
tf.summary.scalar('val/accuracy', acc, step=step)
tf.summary.text('config', 'lr=1e-3, batch=32', step=0)
writer.flush()
# in a custom training loop
for step, (x, y) in enumerate(train_ds):
# ... train step ...
if step % 100 == 0:
with writer.as_default():
tf.summary.scalar('loss', loss, step=step)用 ResNet 迁移学习
特征提取冻结整个预训练骨干网络,只训练新头——最快、最安全的阶段,小数据集或与 ImageNet 相似时理想。骨干网络用 training=False 调用使 BatchNorm 保持预训练统计(小 batch 上 training=True 会更新并破坏特征)。仅头训练时高 LR(1e-3)即可。
import tensorflow as tf
writer = tf.summary.create_file_writer('logs/viz')
with writer.as_default():
tf.summary.image('sample', img[None, ...], step=step)
tf.summary.image('batch', images[:8], step=step, max_outputs=8)
for var in model.trainable_variables:
tf.summary.histogram('weights/' + var.name, var, step=step)
tf.summary.audio('waveform', audio[None, ...], sample_rate=16000, step=step)
writer.flush()微调
第二阶段:解冻顶部层,用 10-100 倍更小的 LR(1e-5)微调以温和适应特征。更改 trainable 后重新编译以重建优化器变量列表。保持早期层冻结,小数据集考虑保持 BatchNorm 冻结。始终在头训练之后再微调。
import tensorflow as tf
# from tensorboard.plugins.hparams import api as hp
# (run: pip install tensorboard-plugin-hparams)
HP_LR = hp.HParam('lr', hp.RealInterval(1e-4, 1e-2))
HP_DROPOUT = hp.HParam('dropout', hp.Discrete([0.1, 0.3, 0.5]))
METRIC_ACC = 'accuracy'
with tf.summary.create_file_writer('logs/hparam_tuning').as_default():
hp.hparams_config(hparams=[HP_LR, HP_DROPOUT], metrics=[hp.Metric(METRIC_ACC)])
for lr in [1e-3, 5e-3]:
for dropout in [0.1, 0.3]:
hparams = {HP_LR: lr, HP_DROPOUT: dropout}
run_name = f"lr-{lr}-drop-{dropout}"
with tf.summary.create_file_writer('logs/hparam_tuning/' + run_name).as_default():
hp.hparams(hparams)
acc = train(hparams)
tf.summary.scalar(METRIC_ACC, acc, step=1)视觉数据管道
image_dataset_from_directory 从类别命名子文件夹读取图像并 yield (图像, 标签) 批——在图像文件夹上开始训练的最快方式。设 image_size 匹配模型输入,label_mode('categorical' one-hot,'int' 稀疏,'binary' 二分类)。加 prefetch 提升性能。
import tensorflow as tf
# profile via callback
tb = tf.keras.callbacks.TensorBoard(
log_dir='logs/profile', profile_batch='10,20',
)
# or use the Profiler API directly
tf.profiler.experimental.start('logs/profile')
# ... run a few batches ...
tf.profiler.experimental.stop()
# in TensorBoard, open the Profile tab to see:
# - op-level time breakdown
# - input pipeline analysis (tf.data bottleneck)
# - memory viewer
# - overview page with recommendationsRNN 文本处理
文本向量化
TextVectorization 在模型图内将原始字符串映射为整数 token ID——adapt() 构建词表。output_sequence_length 填充/截断到固定长度。Embedding 上 mask_zero=True 传播掩码使 RNN 忽略填充——对变长序列很重要。
import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
print('GPUs:', gpus)
# set memory growth (don't grab all GPU memory at start)
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
# or limit GPU memory explicitly
tf.config.set_logical_device_configuration(
gpus[0],
[tf.config.LogicalDeviceConfiguration(memory_limit=4096)],
)
with tf.device('/GPU:1'):
x = tf.random.normal((1000, 1000))
y = tf.matmul(x, x)词嵌入
Embedding 将 token ID 映射为训练中学习的稠密向量——捕获语义相似性。小数据集用对齐词表的预训练 GloVe/word2vec 初始化并冻结(trainable=False)以防止过拟合。大数据集从头训练或微调预训练的通常更好。
import tensorflow as tf
strategy = tf.distribute.OneDeviceStrategy(device='/gpu:0')
strategy = tf.distribute.MirroredStrategy()
print('num replicas:', strategy.num_replicas_in_sync)
# TPU
resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)
with strategy.scope():
model = build_model()
model.compile(optimizer='adam', loss='...', metrics=['accuracy'])
model.fit(train_ds, epochs=10)LSTM 情感分类器
IMDB 情感分析是经典文本分类基线。Embedding、Bidirectional LSTM、Dense 头达约 87%。Bidirectional 让 LSTM 双向读取(免费质量提升)。Dropout 防止小数据集过拟合。SOTA 用 Transformer 编码器替换 LSTM,但 LSTM 仍是强、快、低资源基线。
import tensorflow as tf
strategy = tf.distribute.MirroredStrategy()
print(f'using {strategy.num_replicas_in_sync} GPUs')
global_batch = 64 * strategy.num_replicas_in_sync
train_ds = make_dataset().batch(global_batch)
with strategy.scope():
model = build_model()
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
model.fit(train_ds, epochs=10, validation_data=val_ds)文本生成(字符级)
字符级文本生成从固定上下文预测下一字符。将每序列分为输入(字符 0..n-1)和目标(字符 1..n)——同一序列偏移一——使模型并行学习预测每下一字符。Embedding、LSTM(return_sequences=True)、Dense(vocab_size) 输出 logits。生成时用温度采样。
import tensorflow as tf
try:
resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)
print('TPUs:', strategy.num_replicas_in_sync)
except ValueError:
print('No TPU found')
strategy = tf.distribute.get_strategy()
with strategy.scope():
model = tf.keras.Sequential([
tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
tf.keras.layers.Dense(10, activation='softmax'),
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
model.fit(x_train, y_train, batch_size=1024, epochs=5)带温度的采样
生成每次采样一个字符,将预测反馈为输入。温度在 softmax 前缩放 logits:<1 锐化分布(保守、重复);>1 展平(创意、冒险);1 是模型真实分布。tf.random.categorical 从 logits 采样。
import tensorflow as tf
tf.keras.mixed_precision.set_global_policy('mixed_float16')
# or bfloat16 on Ampere+ / TPU (no loss scaling needed)
tf.keras.mixed_precision.set_global_policy('mixed_bfloat16')
model = tf.keras.Sequential([
tf.keras.layers.Dense(512, activation='relu', input_shape=(784,)),
tf.keras.layers.Dense(10, dtype='float32'),
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
print(model.layers[0].dtype_policy) # <Policy "mixed_float16">Transformer
位置编码
自注意力是置换不变的,所以 Transformer 显式注入位置信息。正弦编码用不同频率的 sin/cos 使每个位置获得唯一、平滑的向量——相对位置是彼此的线性函数。加位置编码前用 sqrt(d_model) 缩放嵌入(Vaswani 配方)。
import tensorflow as tf
augment = tf.keras.Sequential([
tf.keras.layers.RandomFlip('horizontal'),
tf.keras.layers.RandomRotation(0.1),
tf.keras.layers.RandomZoom(0.1),
tf.keras.layers.RandomContrast(0.2),
tf.keras.layers.RandomTranslation(0.1, 0.1),
])
# put augmentation INSIDE the model (training-only, inactive at inference)
model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(224, 224, 3)),
tf.keras.layers.Rescaling(1./255),
augment,
tf.keras.layers.Conv2D(32, 3, activation='relu'),
])
# or apply via dataset.map
ds = ds.map(lambda x, y: (augment(x, training=True), y))多头注意力
MultiHeadAttention 将 d_model 分为 num_heads 子空间,每个内注意力,然后拼接——让 模型并行关注不同关系。自注意力时 query/key/value 传同一张量。attention_mask 阻止位置对:下三角掩码给因果(解码器)注意力。return_attention_scores=True 返回 softmax 权重用于可视化。
from tensorflow.keras.preprocessing.image import ImageDataGenerator
train_datagen = ImageDataGenerator(
rescale=1./255, rotation_range=20,
width_shift_range=0.2, height_shift_range=0.2,
shear_range=0.2, zoom_range=0.2,
horizontal_flip=True, fill_mode='nearest',
validation_split=0.2,
)
train_gen = train_datagen.flow_from_directory(
'data/train', target_size=(224, 224), batch_size=32,
class_mode='binary', subset='training',
)
model.fit(train_gen, epochs=10)Transformer 编码器块
Transformer 编码器块是 MultiHeadAttention + 前馈网络,各用残差连接和 LayerNorm 包裹。Pre-LN(子层前归一化)比原始 Post-LN 训练更稳定,无需 warmup。dff 通常为 d_model 的 4 倍。堆叠 N=6 块达 BERT-base 规模。传 attention_mask 忽略填充 token。
import tensorflow as tf
def augment(image, label):
image = tf.image.random_flip_left_right(image)
image = tf.image.random_brightness(image, max_delta=0.2)
image = tf.image.random_contrast(image, lower=0.8, upper=1.2)
image = tf.image.random_saturation(image, lower=0.8, upper=1.2)
image = tf.image.random_hue(image, max_delta=0.1)
image = tf.clip_by_value(image, 0.0, 1.0)
return image, label
ds = ds.map(augment, num_parallel_calls=tf.data.AUTOTUNE)Transformer 分类器
Transformer 编码器 + 池化头构成强文本分类器。对序列均值池化(或取 [CLS] token)然后 Dense softmax。用 BERT 优化器设置(Adam beta_1=0.9、beta_2=0.98、epsilon=1e-9)和 warmup+cosine LR 调度——transformer 早期对 LR 敏感(故需 warmup)。SOTA 需微调预训练 BERT。
import tensorflow as tf
def mixup(images, labels, alpha=0.2):
lam = tf.random.beta([], alpha, alpha)
batch_size = tf.shape(images)[0]
idx = tf.random.shuffle(tf.range(batch_size))
mixed_images = lam * images + (1 - lam) * tf.gather(images, idx)
mixed_labels = lam * labels + (1 - lam) * tf.gather(labels, idx)
return mixed_images, mixed_labels
ds = ds.batch(32).map(mixup, num_parallel_calls=tf.data.AUTOTUNE)解码器与因果掩码
解码器块有两个注意力层:(1) 对目标的掩码自注意力(因果掩码防止看到未来 token),(2) 交叉注意力(query 来自解码器,key/value 来自编码器输出)。自注意力合并因果和填充掩码。推理时每层缓存 key/value(KV 缓存)。
import tensorflow as tf
import random
def text_augment(text):
words = text.split()
if random.random() < 0.2:
words = [w for w in words if random.random() > 0.2]
if len(words) > 1 and random.random() < 0.2:
i, j = random.sample(range(len(words)), 2)
words[i], words[j] = words[j], words[i]
return ' '.join(words)
def tabular_mixup(x, y, alpha=0.2):
lam = tf.random.beta([], alpha, alpha)
idx = tf.random.shuffle(tf.range(tf.shape(x)[0]))
return lam * x + (1 - lam) * tf.gather(x, idx), lam * y + (1 - lam) * tf.gather(y, idx)
def add_noise(x, std=0.01):
return x + tf.random.normal(tf.shape(x), stddev=std)迁移学习
特征提取(冻结)
特征提取冻结整个预训练骨干网络,只训练新头——最快、最安全的迁移学习阶段,小数据集或与 ImageNet 相似时理想。骨干网络输出特征图;GlobalAveragePooling 展平。骨干网络用 training=False 调用以使 BatchNorm 保持预训练统计。仅头训练时高 LR(1e-3)即可。
import tensorflow as tf
base = tf.keras.applications.ResNet50(
include_top=False, weights='imagenet',
input_shape=(224, 224, 3), pooling='avg',
)
base.trainable = False
inputs = tf.keras.Input(shape=(224, 224, 3))
x = tf.keras.applications.resnet50.preprocess_input(inputs)
x = base(x, training=False)
x = tf.keras.layers.Dense(256, activation='relu')(x)
x = tf.keras.layers.Dropout(0.5)(x)
outputs = tf.keras.layers.Dense(num_classes, activation='softmax')(x)
model = tf.keras.Model(inputs, outputs)两阶段微调
两阶段迁移学习:(1) 冻结骨干网络,只训练头(高 LR);(2) 解冻顶部层,用 10-100 倍更小的 LR(1e-5)微调。头必须先训练——从随机头微调会向骨干网络发送大梯度并破坏预训练特征。更改 trainable 后重新编译。小数据集保持 BatchNorm 冻结。
import tensorflow as tf
base = tf.keras.applications.MobileNetV2(
include_top=False, weights='imagenet', pooling='avg')
base.trainable = False
model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(224, 224, 3)),
tf.keras.applications.mobilenet_v2.preprocess_input,
base,
tf.keras.layers.Dropout(0.3),
tf.keras.layers.Dense(num_classes, activation='softmax'),
])
model.compile(optimizer=tf.keras.optimizers.Adam(1e-3),
loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(train_ds, validation_data=val_ds, epochs=5)预训练文本模型(BERT)
TF Hub 提供带匹配预处理层的预训练 BERT——输入原始字符串,获得上下文嵌入。用小 LR(2e-5 是 BERT 标准)微调 2-4 轮(更多会过拟合)。句子级任务用 pooled_output,token 级任务用 sequence_output。完整 BERT-base 换 L-12_H-768_A-12。
import tensorflow as tf
base = tf.keras.applications.ResNet50(include_top=False, weights='imagenet', pooling='avg')
base.trainable = True
# freeze everything except the last few blocks
for layer in base.layers[:-20]:
layer.trainable = False
model.compile(
optimizer=tf.keras.optimizers.Adam(1e-5), # 10x smaller than head
loss='sparse_categorical_crossentropy', metrics=['accuracy'],
)
model.fit(train_ds, validation_data=val_ds, epochs=5)Hub 层与嵌入
通用句子编码器将句子映射为固定 512 维向量,捕获语义——Keras 即插即用层,适用于文本分类、语义搜索或聚类,训练数据很少。冻结(trainable=False)以提速。句子相似度是嵌入的余弦。生产用 sentence-transformers 更好,但 USE 是最快 TF 原生选项。
import tensorflow as tf
base = tf.keras.applications.EfficientNetB0(
include_top=False, weights='imagenet', pooling='avg')
base.trainable = False
for layer in base.layers:
layer.trainable = 'block7' in layer.name
# freeze BatchNorm specifically (recommended during fine-tuning)
for layer in base.layers:
if isinstance(layer, tf.keras.layers.BatchNormalization):
layer.trainable = False
n_train = sum(tf.keras.backend.count_params(w) for w in model.trainable_weights)
n_total = sum(tf.keras.backend.count_params(w) for w in model.weights)
print(f"trainable: {n_train:,} / {n_total:,}")知识蒸馏
知识蒸馏训练小学生模型模仿大教师模型的软输出(错误类概率中的暗知识)。温度 T 软化分布使学生学习教师的类间相似性。alpha 平衡硬标签(真实类)和软标签(教师输出)。学生以教师一小部分的大小/速度成本接近教师精度。
import tensorflow as tf
base = tf.keras.applications.ResNet50(include_top=False, weights='imagenet')
base.trainable = False
inputs = tf.keras.Input(shape=(224, 224, 3))
x = tf.keras.applications.resnet50.preprocess_input(inputs)
features = base(x, training=False)
features = tf.keras.layers.GlobalAveragePooling2D()(features)
cls_out = tf.keras.layers.Dense(num_classes, activation='softmax', name='cls')(features)
box_out = tf.keras.layers.Dense(4, name='box')(features)
attr_out = tf.keras.layers.Dense(num_attr, activation='sigmoid', name='attr')(features)
model = tf.keras.Model(inputs, [cls_out, box_out, attr_out])
model.compile(
optimizer='adam',
loss={'cls': 'sparse_categorical_crossentropy', 'box': 'mse',
'attr': 'binary_crossentropy'},
loss_weights={'cls': 1.0, 'box': 5.0, 'attr': 0.5},
)模型保存与加载
Keras .keras 格式
.keras 格式(TF 2.13+)是推荐的单文件格式——在一个 zip 中保存架构、权重、优化器状态和自定义对象。save_weights/load_weights 仅存权重(用于冻结或迁移到不同架构)。to_json 仅序列化架构。自定义层用 @tf.keras.utils.register_keras_serializable() 注册以使 load_model 可重建。
import tensorflow as tf
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
with open('model.tflite', 'wb') as f:
f.write(tflite_model)
# with quantization (8-bit, ~4x smaller, faster on CPU)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# float16 quantization (GPU-friendly, 2x smaller)
converter.target_spec.supported_types = [tf.float16]
# full integer quantization (needs representative dataset)
converter.representative_dataset = rep_data
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
tflite_quant = converter.convert()SavedModel 格式
SavedModel 是 TF 的生产就绪格式——包含计算图、变量和签名的目录。TF Serving、TFLite 转换器和 TFX 都消费它。model.export() 创建它(TF 2.13+;之前用 model.save('dir/'))。serving_default 签名是标准推理入口。部署用此格式;训练检查点用 .keras。
import tensorflow as tf
import numpy as np
interpreter = tf.lite.Interpreter(model_path='model.tflite')
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
print(input_details)
# [{'name': 'input', 'shape': (1, 224, 224, 3), 'dtype': tf.float32, ...}]
input_data = np.expand_dims(img, axis=0).astype(np.float32)
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()
output = interpreter.get_tensor(output_details[0]['index'])
print(output.shape, output.argmax())检查点与权重
ModelCheckpoint 回调在 fit 期间保存权重。自定义训练循环用 tf.train.Checkpoint + CheckpointManager 获得更细控制:保存模型、优化器、步数计数器,任何东西——max_to_keep 轮换旧检查点。Checkpoint 显式存储对象图(不像 HDF5 按变量名存储),所以适用于子类化模型和自定义循环。
# save model in SavedModel format with versioning
# model.export('serving/1') # Keras 3
tf.saved_model.save(model, 'serving/1')
# run TF Serving via Docker
# docker run -p 8501:8501 --name tfserving \
# -v $(pwd)/serving:/models/my_model \
# -e MODEL_NAME=my_model \
# tensorflow/serving
# query via REST
import json, requests
data = json.dumps({'instances': x_test[:3].tolist()})
resp = requests.post(
'http://localhost:8501/v1/models/my_model:predict',
data=data, headers={'content-type': 'application/json'})
print(resp.json()['predictions'])保存自定义对象
要使 load_model 重建自定义层,用 @register_keras_serializable 注册(自动发现)或传 custom_objects={'LayerName': LayerClass} 给 load_model。层必须实现 get_config() 返回构造参数。否则加载抛 'Unknown layer' 错误。这也适用于自定义损失、指标和优化器。
import tensorflow as tf
# (requires: pip install tensorflow-model-optimization)
import tensorflow_model_optimization as tfmot
model = tf.keras.applications.MobileNetV2(weights='imagenet')
annot_model = tfmot.quantization.keras.quantize_model(model)
annot_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
annot_model.fit(train_ds, epochs=3)
converter = tf.lite.TFLiteConverter.from_keras_model(annot_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_qat = converter.convert()导出用于推理
生产部署时,用带显式 input_signature 的 @tf.function 将模型包装在 tf.Module 中——这固定输入形状/dtype 使 TF Serving 和 TFLite 知道接口。命名输出张量(dict 键)以清晰。tf.saved_model.save 配 signatures 使函数作为 serving_default 可用。这是将模型交付生产的标准方式。
# convert to TF.js format
# pip install tensorflowjs
# tensorflowjs_converter --input_format=keras model.keras tfjs_model/
# in JavaScript:
# import * as tf from '@tensorflow/tfjs';
# const model = await tf.loadLayersModel('tfjs_model/model.json');
# const out = model.predict(tf.tensor4d(img, [1, 224, 224, 3]));
# OR convert SavedModel -> TF.js
# tensorflowjs_converter \
# --input_format=tf_saved_model \
# --output_format=tfjs_graph_model \
# saved_model/ tfjs_model/
# TF.js layers model: full Keras-like API in the browser
# TF.js graph model: lower-level, supports more ops检查点回调
长时间训练运行时全面检查点:最佳模型 .keras(用于部署)、周期性权重(用于恢复)、BackupAndRestore(用于崩溃恢复)。save_freq='epoch' 每轮保存;用整数 N 每 N 批保存。{epoch:02d} 占位符在文件名中嵌入轮次号。BackupAndRestore 对不可靠硬件或长分布式作业必需。
import tensorflow as tf
# knowledge distillation: train a small student to mimic a large teacher
def distill_loss(teacher_logits, student_logits, temperature=3.0):
soft_targets = tf.nn.softmax(teacher_logits / temperature)
soft_probs = tf.nn.log_softmax(student_logits / temperature)
return tf.reduce_mean(
tf.keras.losses.categorical_crossentropy(soft_targets, soft_probs, from_logits=True))
# XLA compilation (faster inference)
@tf.function(jit_compile=True)
def predict(x):
return model(x)
# stack optimizations: prune + quantize + distill + XLA
# always benchmark latency and accuracy before/after each stepTensorBoard 可视化
TensorBoard 基础设置
TensorBoard 记录损失/指标、权重直方图、计算图和性能分析数据用于浏览器实时可视化。histogram_freq=1 每轮追踪权重分布。profile_batch 分析一批的计算时间。用 `tensorboard --logdir <dir>` 启动,打开 localhost:6006。子目录中多次运行显示为叠加曲线以便比较。
import tensorflow as tf
optimizer = tf.keras.optimizers.Adam(1e-3)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
train_loss = tf.keras.metrics.Mean(name='train_loss')
train_acc = tf.keras.metrics.SparseCategoricalAccuracy(name='train_acc')
@tf.function
def train_step(x, y):
with tf.GradientTape() as tape:
logits = model(x, training=True)
loss = loss_fn(y, logits)
loss += sum(model.losses)
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
train_loss.update_state(loss)
train_acc.update_state(y, logits)
return loss
for epoch in range(epochs):
for x, y in train_ds:
train_step(x, y)
print(f"epoch {epoch}: loss={train_loss.result():.4f} acc={train_acc.result():.4f}")
train_loss.reset_state(); train_acc.reset_state()自定义标量与图像
tf.summary.* 将自定义数据记录到 TensorBoard:标量(指标、LR、梯度范数)、图像(增强样本、激活、Grad-CAM)、直方图(权重分布)和文本(配置、超参)。创建文件写入器并用 writer.as_default() 包裹。调用 writer.flush() 强制写入。这是可视化标准 fit 指标以外内容的方式。
import tensorflow as tf
val_loss = tf.keras.metrics.Mean(name='val_loss')
val_acc = tf.keras.metrics.SparseCategoricalAccuracy(name='val_acc')
@tf.function
def test_step(x, y):
logits = model(x, training=False)
loss = loss_fn(y, logits)
val_loss.update_state(loss)
val_acc.update_state(y, logits)
for epoch in range(epochs):
for x, y in train_ds:
train_step(x, y)
for x, y in val_ds:
test_step(x, y)
print(f"epoch {epoch}: val_loss={val_loss.result():.4f} val_acc={val_acc.result():.4f}")
train_loss.reset_state(); train_acc.reset_state()
val_loss.reset_state(); val_acc.reset_state()性能分析
Profiler 分解训练时间去向:运算级计时、内存使用、设备利用率和输入管道效率。最常见发现是输入管道瓶颈(GPU 饥饿而 CPU 解码图像)——用 prefetch 和 num_parallel_calls 修复。Profile 标签还建议具体优化。定期分析而非每次运行,因为有额外开销。
import tensorflow as tf
strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
model = build_model()
optimizer = tf.keras.optimizers.Adam(1e-3)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
train_ds_dist = strategy.experimental_distribute_dataset(train_ds)
@tf.function
def distributed_train_step(x, y):
def step_fn(x, y):
with tf.GradientTape() as tape:
logits = model(x, training=True)
loss = loss_fn(y, logits)
loss += sum(model.losses)
loss /= strategy.num_replicas_in_sync
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
return loss
per_replica_loss = strategy.run(step_fn, args=(x, y))
return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_loss, axis=None)
for epoch in range(epochs):
for x, y in train_ds_dist:
distributed_train_step(x, y)HParams 调参面板
HParams 面板记录超参运行使 TensorBoard 可显示平行坐标、散点图和表格视图比较不同配置的准确率。用 hp.HParam 定义搜索空间,用 hp.hparams() 记录每次运行的超参+指标。当已用 TensorBoard 时,这是专用 HPO 工具(KerasTuner、Optuna)的轻量替代。
import tensorflow as tf
generator = build_generator()
discriminator = build_discriminator()
g_opt = tf.keras.optimizers.Adam(1e-4, beta_1=0.5)
d_opt = tf.keras.optimizers.Adam(1e-4, beta_1=0.5)
bce = tf.keras.losses.BinaryCrossentropy(from_logits=True)
@tf.function
def train_gan_step(real_images):
batch = tf.shape(real_images)[0]
noise = tf.random.normal((batch, 100))
with tf.GradientTape() as g_tape, tf.GradientTape() as d_tape:
fake = generator(noise, training=True)
real_logits = discriminator(real_images, training=True)
fake_logits = discriminator(fake, training=True)
d_loss = bce(tf.ones_like(real_logits), real_logits) + \
bce(tf.zeros_like(fake_logits), fake_logits)
g_loss = bce(tf.ones_like(fake_logits), fake_logits)
g_grads = g_tape.gradient(g_loss, generator.trainable_variables)
d_grads = d_tape.gradient(d_loss, discriminator.trainable_variables)
g_opt.apply_gradients(zip(g_grads, generator.trainable_variables))
d_opt.apply_gradients(zip(d_grads, discriminator.trainable_variables))
return g_loss, d_loss嵌入投影器
Projector 通过 PCA 或 t-SNE 在 2D/3D 中可视化高维嵌入(词向量、图像特征),揭示聚类和类比(king - man + woman ≈ queen)。将嵌入矩阵保存为检查点变量,metadata.tsv 将行映射到标签。调试 NLP 模型极有用——可看到相似词是否聚类在一起。
import tensorflow as tf
accum_steps = 4
optimizer = tf.keras.optimizers.Adam(1e-3)
@tf.function
def accumulated_train_step(x, y):
with tf.GradientTape() as tape:
logits = model(x, training=True)
loss = loss_fn(y, logits) / accum_steps
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
return loss
# simple accumulation via averaging
for batch_idx, (x, y) in enumerate(train_ds):
accumulated_train_step(x, y)
# equivalent effective batch = accum_steps * batch_size
# (Adam moments are updated every step, so this is an approximation)GPU/TPU 与部署
GPU 内存与策略
默认 TF 启动时分配所有 GPU 内存。set_memory_growth=True 增量分配(使其他进程可共享 GPU)。测试时用 set_logical_device_configuration 将一个物理 GPU 分成多个逻辑设备。内存限制在训练和服务共享 GPU 时有用。多 GPU 机器用 CUDA_VISIBLE_DEVICES 限制 TF 可见的 GPU。
import tensorflow as tf
@tf.function
def train_step(x, y):
with tf.GradientTape() as tape:
preds = model(x, training=True)
loss = loss_fn(y, preds)
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
return loss
# first call traces the graph (slow), subsequent calls reuse it (fast)
train_step(x_batch, y_batch)
# experimental_relax_shapes allows retracing only when needed
@tf.function(reduce_retracing=True)
def fn(x): return model(x)
# input_signature pins the signature so no retracing
@tf.function(input_signature=[tf.TensorSpec([None, 784], tf.float32)])
def predict(x): return model(x)MirroredStrategy(多 GPU)
MirroredStrategy 做同步数据并行训练:变量在每个 GPU 上复制,每副本处理 batch 的一片,梯度 all-reduce(默认 NCCL)并平均。全局 batch size 按 num_replicas_in_sync 缩放以保持每 GPU batch 相同。在 strategy.scope() 内构建模型。多机用 MultiWorkerMirroredStrategy;TPU 用 TPUStrategy。
import tensorflow as tf
# JIT-compile a function with XLA (fuses ops for speed)
@tf.function(jit_compile=True)
def train_step(x, y):
with tf.GradientTape() as tape:
loss = loss_fn(y, model(x, training=True))
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
return loss
# enable XLA globally
tf.config.optimizer.set_jit(True)
# XLA fuses element-wise ops, reduces kernel launches,
# and can give 1.5-3x speedup on GPU/TPU
# especially effective for transformer / attention heavy modelsTPU 策略
TPU(Tensor Processing Unit)是 Google 的 ML 加速器——大型矩阵乘法快。用 TPUStrategy:通过 TPUClusterResolver 连接,初始化 TPU 系统,然后在 scope() 内构建模型。TPU 仅从 Google Cloud Storage(GCS)或内存中 tf.data 高效读取数据,非本地磁盘。TPU 擅长大 batch 和 bfloat16 训练。
import tensorflow as tf
# the canonical high-performance pipeline
ds = (tf.data.Dataset.from_tensor_slices((x, y))
.cache() # cache after loading
.shuffle(10000)
.map(augment, num_parallel_calls=tf.data.AUTOTUNE)
.batch(32, drop_remainder=True)
.prefetch(tf.data.AUTOTUNE))
# parallel file reading for TFRecord
files = tf.data.Dataset.list_files('data-*.tfrecord')
ds = files.interleave(
tf.data.TFRecordDataset,
cycle_length=4, num_parallel_calls=tf.data.AUTOTUNE)
# options for autotuning
options = tf.data.Options()
options.autotune.enabled = True
ds = ds.with_options(options)TF Serving 部署
TF Serving 是生产级模型服务器:加载 SavedModel,暴露 REST(端口 8501)和 gRPC(端口 8500)端点,处理批处理、版本管理(数字子目录)和新版本热替换不停机。Docker 是最简部署。高吞吐时 gRPC 比 REST 延迟更低。TF Serving 自动服务最高版本号,支持金丝雀发布。
import tensorflow as tf
# mixed precision (float16 compute, float32 master)
tf.keras.mixed_precision.set_global_policy('mixed_float16')
# enable TF32 on Ampere+ (big matmul speedup, slight precision loss)
tf.config.experimental.enable_tensor_float_32_execution(True)
# TF32 uses 19-bit mantissa for matmul on A100/H100
# it's transparent — no code changes, no loss scaling
# gives ~3x matmul speedup with minimal accuracy impact
# check policy
print(tf.keras.mixed_precision.global_policy())
print('TF32:', tf.config.experimental.tensor_float_32_execution_enabled())TF Lite(移动端与边缘)
TF Lite 将模型部署到移动端(Android/iOS)、边缘设备(树莓派)和微控制器。转换器将 SavedModel 转为扁平 .tflite 缓冲区;Optimize.DEFAULT 应用动态范围量化(约 4 倍小,精度损失极小)。完整 int8 量化(Edge TPU 上最快)需提供 representative_dataset。Interpreter API 在设备上用相同输入/输出张量接口推理。
import tensorflow as tf
# enable grappler (default graph optimizer)
tf.config.optimizer.set_jit(True) # XLA
tf.config.optimizer.set_experimental_options({
'disable_model_pruning': False,
'disable_meta_optimizer': False,
'layout_optimizer': True, # optimize data layout (NCHW vs NHWC)
'constant_folding': True, # fold constants at build time
'shape_optimization': True,
'remapping': True, # fuse compatible ops
'arithmetic_optimization': True,
'loop_optimization': True,
'dependency_optimization': True,
'function_optimization': True,
'debug_stripper': True, # remove Assert/CheckNumerics
})
# all on by default — only set explicitly to disable specific onesTF.js(浏览器部署)
TF.js 用 WebGL/WebGPU 加速在浏览器或 Node.js 中直接运行模型。转换器将 SavedModel/Keras 模型转为 model.json + 二进制权重分片。用 tf.loadGraphModel 或 tf.loadLayersModel 加载。适合隐私保护推理(数据不离开设备)、交互式演示和离线优先应用。量化选项(uint8/uint16)减小下载大小。
import tensorflow as tf
# AutoGraph converts Python if/for/while into tf.cond / tf.while_loop
@tf.function
def fn(x):
if x > 0: # AutoGraph converts this
return x * 2
else:
return x
@tf.function
def sum_loop(n):
total = tf.constant(0)
for i in tf.range(n): # AutoGraph converts to tf.while_loop
total += i
return total
# Python-side (compile-time) loop vs TF-side (runtime) loop
@tf.function
def train_epoch(dataset):
for x, y in dataset: # Python iteration over dataset (fine)
train_step(x, y)
# avoid data-dependent Python control flow
# if x.shape[0] > 10: ... # OK (static shape)
# if tf.reduce_sum(x) > 0: ... # forces retracing相关 TensorFlow 代码片段
Copy-paste ready code for common tasks.
Tensor Basics
Create and operate on TensorFlow tensors.
Keras Model
Build models with Sequential and the functional API.
Layers
Use core layers and build a custom one.
Compile and Train
Compile, fit, and evaluate a Keras model.
Custom Training Loop
Step through batches with GradientTape.
Callbacks
Monitor and control training with callbacks.
Save and Load
Persist models in SavedModel and Keras formats.
Data Pipeline
Build efficient input pipelines with tf.data.
这篇内容对您有帮助吗？