3、TensorFlow基础快速回顾：TensorFlow 2.x核心API、Eager Execution、Keras高层接口、SavedModel格式

说实话，很多从PyTorch转过来的朋友，一听到TensorFlow就头疼。我完全理解。早年的TF 1.x确实有点反人类——静态图、Session、placeholder，写个简单的线性回归都要绕半天。

但TF 2.x完全不一样了。它几乎把PyTorch那套「动态图」的体验搬了过来，同时又保留了生产部署上的优势。今天我们就快速过一遍TF 2.x的核心内容，为后面跟JAX的融合打好基础。

3.1 Eager Execution：终于不用再「编译」了

Eager Execution，说白了就是「立即执行」。你写一行代码，它立刻算出一个结果，而不是先搭个计算图再等Session.run()。

我记得第一次用TF 1.x时，写了个简单的矩阵乘法，结果打印出来是个Tensor对象，不是数值。当时我就懵了。现在好了，TF 2.x默认开启Eager模式，代码写起来跟NumPy一样自然。

import tensorflow as tf

# Eager Execution 默认开启
a = tf.constant([[1, 2], [3, 4]])
b = tf.constant([[5, 6], [7, 8]])
c = tf.matmul(a, b)

print(c.numpy())  # 直接拿到数值，不用Session
# 输出：
# [[19 22]
#  [43 50]]

为什么会这样？因为TF 2.x把每个操作都变成了「立即执行」的算子。你调用tf.matmul，它当场就做矩阵乘法，返回一个EagerTensor。你再调用.numpy()，就能拿到NumPy数组。

3.2 Keras高层接口：写模型就该这么简单

Keras现在已经是TF的官方高层API了。你想想看，以前用TF 1.x搭个CNN，要写变量作用域、占位符、损失函数、优化器……一堆模板代码。现在用Keras，几行搞定。

Keras提供了三种建模方式，我按推荐程度排个序：

方式	适用场景	我的评价
Sequential（顺序模型）	简单堆叠，没有分支	入门最快，但灵活性差
Functional API（函数式API）	多输入、多输出、共享层	我最常用的方式，平衡了简洁和灵活
Model Subclassing（子类化）	完全自定义，研究用	跟PyTorch的nn.Module很像，自由度最高

举个例子，用Functional API搭一个简单的残差块：

from tensorflow.keras import layers, Model

def residual_block(x, filters):
    shortcut = x
    x = layers.Conv2D(filters, 3, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Conv2D(filters, 3, padding='same')(x)
    x = layers.BatchNormalization()(x)
    # 残差连接
    x = layers.add([x, shortcut])
    x = layers.ReLU()(x)
    return x

inputs = layers.Input(shape=(32, 32, 3))
x = layers.Conv2D(64, 3, padding='same')(inputs)
x = residual_block(x, 64)
x = layers.GlobalAveragePooling2D()(x)
outputs = layers.Dense(10, activation='softmax')(x)

model = Model(inputs, outputs)
model.summary()

嗯，这里要注意：Functional API的层之间是「有向无环图」结构。你不能在里面写循环或者条件分支。如果需要那种动态控制流，就得用Model Subclassing。

3.3 SavedModel格式：部署的「通用语言」

SavedModel是TF官方推荐的模型序列化格式。它不像HDF5那样只保存权重，而是把整个计算图、权重、签名函数都打包在一起。

我为什么说它重要？因为后面跟JAX融合时，我们经常需要把JAX模型转成SavedModel，然后丢到TF Serving或者TFLite里跑。说白了，SavedModel就是TF生态的「通用接口」。

保存和加载的代码很简单：

# 保存
model.save('my_model', save_format='tf')  # 默认就是SavedModel

# 加载
loaded_model = tf.keras.models.load_model('my_model')

# 验证一下
import numpy as np
test_input = np.random.randn(1, 32, 32, 3).astype(np.float32)
output = loaded_model(test_input)
print(output.shape)  # (1, 10)

SavedModel的目录结构是这样的：

my_model/
├── assets/
├── variables/
│   ├── variables.data-00000-of-00001
│   └── variables.index
├── saved_model.pb   # 计算图定义
└── keras_metadata.pb  # Keras特有信息

你想想看，这个结构多清晰。.pb文件存的是图结构，variables文件夹存的是权重。不管你是用Python、C++还是Java加载，都能还原出完整的模型。

3.4 本章知识体系

下面这张图，帮你把TF 2.x的核心脉络理清楚：

这张图把TF 2.x的三大核心串起来了。Eager Execution让你写代码像写Python脚本一样自然；Keras让你搭模型像搭积木一样简单；SavedModel让你部署模型像发快递一样标准。

嗯，到这里，TF 2.x的基础回顾就差不多了。你可能会问：「这些跟JAX有什么关系？」别急，下一章我们就要开始动手了——用JAX写一个简单的神经网络，然后跟TF的Keras模型做对比。你会发现，两者在底层其实有很多相通的地方。