做了几年深度学习，我发现自己大部分的调试时间其实不在网络结构上，而是在张量操作上。框架装好、数据到位、模型搭起来之后，真正拦路的反而全是shape对不上、维度交换错了、广播不按预期走这种基础问题。TensorFlow里的张量操作，从最简单的创建、切片，到高性能计算里的算子融合、数据流水线，是一条值得从头梳理清楚的主线。这篇文章就围绕TensorFlow张量操作展开，从基本概念和常用操作讲起，再进入tf.function、XLA、混合精度这些性能优化手段，最后用一个实际案例演示怎么一步步榨干显存和算力。适合刚接触TensorFlow、被各种rank和shape绕晕的新手，也适合已经能跑通模型、但训练和推理速度始终上不去的进阶开发者。
1. 张量是什么：先搞清楚Tensor的三重身份
1.1 张量的三重身份
讲操作之前，得先把张量这个东西想清楚。很多初学者把Tensor简单理解成“多维数组”，这个说法对，但不够。在TensorFlow里，一个张量其实有三重身份，理解这三重身份，后面所有的操作你都能立刻判断它会不会引发性能问题或者报错。
第一重身份是数据结构。一个Tensor由两部分组成：形状（shape）和数据类型（dtype）。shape描述每个维度的大小，dtype描述元素类型，比如tf.float32、tf.int64、tf.bool。这里有一个很容易被忽略的点：张量本身不存储数据，它是数据的“视图”。它存储的是数据在底层连续内存中的偏移和步长信息。所以对张量做reshape，往往只是在修改视图的元数据，而不是真的搬运内存数据。
第二重身份是计算图中的节点。虽然TensorFlow 2.x默认采用动态图模式，看起来像普通Python函数一样一行行执行，但底层仍然有一个计算图在调度。每个算子都是一个节点，张量是节点之间的边。这个视角对性能至关重要，因为图模式允许框架做算子融合、并发调度、内存复用，这是后面讲高性能计算的基础。
第三重身份是设备上的物理存储。一个Tensor要么在CPU内存里，要么在GPU显存里。两种设备之间的张量复制开销非常大，一次跨越PCIe的数据拷贝可能比一次小规模矩阵乘法还慢。
我见过不少人排查性能问题时，明明代码逻辑没问题，最后发现是CPU和显存之间的数据交换太频繁，导致GPU一直在空转。这就是因为没意识到张量还带着“设备”这个属性。
1.2 张量的创建工具箱
创建张量的方式很多，搞懂每种方式的适用场景，能少写很多无用代码。
import tensorflow as tf # 常量：不可变，适合保存固定参数，比如归一化均值 a = tf.constant(3.14) b = tf.constant([1, 2, 3], dtype=tf.int32) # 从Python列表或NumPy数组转换 c = tf.convert_to_tensor([[1.0, 2.0], [3.0, 4.0]]) # 变量：可训练，适合保存模型权重 v = tf.Variable(tf.zeros([2, 2])) # 随机初始化，注意每个生成函数都有固定的分布形态 r1 = tf.random.normal([32, 128]) # 正态分布 r2 = tf.random.uniform([32, 128]) # 均匀分布 # 特殊结构张量 z = tf.zeros([5, 10]) o = tf.ones([3, 4]) e = tf.eye(5) # 单位矩阵 seq = tf.linspace(0.0, 1.0, 11) # 等差数列，注意包含端点这里有几个实操经验可以分享：
tf.random.normal不指定seed时，每次执行结果都不一样。在调试阶段，如果不固定随机种子，同样的代码每次跑出来的指标都不同，非常容易让人误判模型效果。调试期给tf.random.set_seed(42)加上，测性能的时候再关掉。
tf.Variable(initial_value=...)的初始值参数是强制的，直接tf.Variable(0.1)会报错。这是很多从PyTorch转过来的朋友容易踩的第一个坑。
tf.constant创建超大张量时要谨慎。它虽然只是保存数据，但如果一次性创建一个几万乘几万的大矩阵，内存立刻就会被打满。正确的做法是先用tf.zeros或tf.random配合后续的tf.tensor_scatter_nd_update分段填充，不要试图在初始化时就构造一个巨大的全量常量。
提示：查看张量信息时，直接打印
print(x)看到的是内容的简化展示，推荐用tf.print(x)或者print(x.shape, x.dtype, x.device)，后者能快速确认设备位置。
2. 形状操作：让数据在层次间自由变换
2.1 reshape、squeeze、expand_dims 各自做什么
形状操作是张量操作里最频繁也最混乱的一块。核心原因在于：很多操作看起来都是“改形状”，但底层做的事情完全不同。
reshape是逻辑重排。它的规则是：把原张量按行优先顺序线性化，再按新形状重新分段。举个生活化的例子：把一本96页的书重新按每章32页装订，你并没有改变书的印刷顺序，只是改变了分组方式。因此reshape严格不改变元素总数，tf.reshape(x, [3, -1])里的-1是自动推算维度，表示“剩下所有元素”。
x = tf.random.normal([4, 6]) y = tf.reshape(x, [3, 8]) # 合法，4*6=3*8 z = tf.reshape(x, [3, -1]) # 自动变成 [3, 8] # 常见错误示例：tf.reshape(x, [4, 4]) 会报错squeeze是去掉维度大小为1的轴。典型场景是CNN中处理批量维度：模型输出[1, 224, 224, 3]，要知道，这个“1”如果没有实际意义，把它留着会干扰后续的矩阵乘法。tf.squeeze(x, axis=0)明确指定去掉batch维，比不指定axis全部去掉更安全，因为有时候维度1是语义上需要保留的。
expand_dims是在指定位置插入一个大小为1的新维度。这个操作在广播场景里几乎每天都要用。比如有一个形状[4]的向量，你想让每个样本分别加上这个向量，直接加会报错；先tf.expand_dims(v, axis=0)变成[1, 4]，就能和[8, 4]的数据广播相加了。
x = tf.random.normal([8, 224, 224, 3]) sq = tf.squeeze(x, axis=0) # 从 [1,224,224,3] 变成 [224,224,3] xd = tf.expand_dims(sq, axis=0) # 把 batch 维加回来reshape和expand_dims有一个关键区别：reshape可能重新排列数据，而expand_dims永远不会，它只是在已有的维度间插入一个“空壳”。所以expand_dims速度极快，开销几乎为零。能用到expand_dims的地方，就不要用reshape去凑形状。
2.2 transpose 为什么比重塑更“重”
transpose操作要单独拿出来讲，因为它是“重操作”的典型代表。
reshape和expand_dims在底层只是改变张量的元数据，不搬动内存数据，速度很快。transpose不同，它要真正交换轴的顺序。对应到内存上，就是数据的物理位置全部打乱重排。这个开销随张量规模线性增长，而且还会破坏内存连续性，导致后续算子访问缓存时性能下降。
x = tf.random.normal([10, 224, 224, 3]) # NHWC 布局 xt = tf.transpose(x, perm=[0, 3, 1, 2]) # 变成 NCHW 布局这里有个经典优化场景：GPU上很多卷积算子对NCHW布局计算效率更高，而CPU上的很多图像处理逻辑喜欢NHWC。在做跨平台部署时，经常需要在两者之间转换。我的建议是：整个流程中尽量固定一种布局，只在必要位置做一次transpose，不要来回切。反复transpose不仅浪费算力，而且切来切去还会让模型结构变得难以维护。
另外，transpose之后的张量视为“非连续”的，如果你在它上面做reshape，TensorFlow往往会被迫先做一次数据拷贝把它变成连续张量，再执行后续操作。所以尽量不要在transpose之后立即reshape。
正确姿势是先计算好最终需要的形状，干完该干的事，最后一次性完成重整形。
2.3 索引、切片和按需取数
张量的索引切片和NumPy的规则很接近，但TensorFlow张量是不可变的，不能直接tensor[idx] = value。
x = tf.random.normal([8, 16, 16, 3]) x0 = x[0] # 取第一张图 [16,16,3] slice1 = x[0:6:2] # 步长为2的切片 [4,16,16,3] cut = x[:, :, :, 0] # 取所有图像的R通道 [8,16,16] # 更高阶的取数 idx = tf.constant([0, 3, 5]) g = tf.gather(x, indices=idx, axis=0) # 按索引取 [3,16,16,3] mask = tf.constant([True, False, True, True, True, True, True, True]) bm = tf.boolean_mask(x, mask, axis=0) # 按布尔掩码取 [7,16,16,3]这里有个很实用的组合技巧：tf.gather配合tf.argsort实现按某一列排序后的重排，常用于榜单类数据的处理。还有tf.gather_nd可以一次取多个位置的元素，比如从特征图里按坐标点采样，比逐点索引再拼接快很多。
如果确实需要修改张量中的某个区域的数值，不能用x[0] = 1这种写法，要用tf.tensor_scatter_nd_update：
x = tf.zeros([4, 4]) updates = tf.constant([1.0, 2.0, 3.0]) indices = tf.constant([[0, 0], [1, 1], [2, 2]]) result = tf.tensor_scatter_nd_update(x, indices, updates)这个操作在实现一些稀疏更新的自定义逻辑时非常有用，但注意频繁使用会带来额外的内存分配开销，如果可以整块更新，优先考虑tf.where这种向量化方式。
3. 张量运算：从广播机制到矩阵乘法
3.1 广播机制：规则、例子和翻车现场
广播（broadcast）是TensorFlow里最反直觉，也最常出错的机制之一。简单说，当两个张量形状不完全相同但兼容时，TF会自动扩展尺寸较小的张量，让它们可以按元素计算。规则只有一条核心逻辑：从最后一个维度开始向前对比，两个维度要么相等，要么其中一个为1，要么其中一方缺失。只要有一个维度不符合，就报Incompatible shapes错误。
a = tf.ones([4, 3]) b = tf.ones([3]) # 缺失第一维，视为 [1,3]，广播到 [4,3] c = a * b # 结果 [4,3] # 合法但容易误解的广播 m = tf.ones([4, 1]) n = tf.ones([1, 3]) o = m * n # 结果是 [4,3]，每一个元素都等于1 # 典型的报错场景 a = tf.ones([4, 2]) b = tf.ones([3]) bad = a * b # InvalidArgumentError生活化类比：广播相当于把一张海报按人数复印分发，每个人手里都拿到同一张海报。关键点是“同一张”而不是“重新设计一张”。所以广播对显存的消耗远小于先手动复制再计算。
我见过最多的翻车现场，发生在某个中间维度尺寸不一致时。比如两个特征矩阵形状分别是[32, 128]和[32, 256]，你想对它们做逐元素计算，广播会立刻报错。这时候不要盲目去reshape，先想清楚你真正想要的对齐方式，是用expand_dims + broadcast_to显式对齐，还是用矩阵乘法的维度匹配规则。这两种方案的语义完全不同。
注意：广播机制不是万能的。它在某些情况下会掩盖逻辑错误。比如一个形状错误的张量，因为另一个张量恰好有维度1，广播就“成功”了，结果看起来数值都对，但语义已经歪了。建议在关键计算之前，用
tf.debugging.assert_equal或check_shape来断言形状。
3.2 归约操作：axis参数的正确理解方式
tf.reduce_sum、tf.reduce_mean这类归约操作，是张量从高维向低维收敛的主要手段。它们最核心的参数就是axis，这个参数理解错了，代码就全错了。
axis指的是你要沿着哪个维度“压扁”。举例：一个形状[2, 3, 4]的张量，可以理解为2个批次、每个批次3行、每行4个元素。
x = tf.random.normal([2, 3, 4]) # axis=0：把2个批次合并，结果是 [3,4] s0 = tf.reduce_sum(x, axis=0) # axis=-1：等价于 axis=2，把每行的4个元素求和，结果是 [2,3] s1 = tf.reduce_sum(x, axis=-1) # axis=[1,2]：同时压缩中间两维，结果是 [2] s2 = tf.reduce_sum(x, axis=[1, 2]) # keepdims 的作用 s3 = tf.reduce_sum(x, axis=-1, keepdims=True) # 结果是 [2,3,1]keepdims=True这个参数很关键。很多初学者在归约后直接拿去和另一个高维张量做广播运算，结果广播规则完全对不上。加上keepdims，归约后保留的维度1就能自动参与广播，省去手动expand_dims的步骤。
我在实际项目里发现一个规律：训练分类任务算每个样本的loss均值、做注意力机制里的softmax归一化、以及实现批归一化时，reduce_mean都少不了。这些场景下，axis的选择直接决定结果是否合理。
有一个容易绕晕的点：做tf.reduce_max(x, axis=0)和tf.argmax(x, axis=0)是不同的。前者返回最大值本身，后者返回最大值的索引。两者配合使用时，要记住argmax的结果形状比原张量低一维，而reduce_max加keepdims后形状保持原来维度。很多人在特征聚合时报维度错，就是没分清这两个操作。
3.3 乘法三兄弟：逐元素、矩阵乘与einsum
TensorFlow里“乘法”有好几种，功能完全不同，新手最容易混淆。
逐元素乘法（*或tf.multiply）：对应位置相乘，形状必须满足广播规则。它适合做缩放、掩码、逐像素加权。
矩阵乘法（tf.matmul或@）：要求左矩阵第二维等于右矩阵第一维，自动处理batch维度。适合做线性层、特征变换。批量矩阵乘法tf.matmul支持两个形状为[batch, m, n]和[batch, n, k]的张量做批量乘法，但两个batch维度必须一致，不一致时先广播再乘。
einsum：一种用爱因斯坦求和约定表达的张量运算。它把各种乘法统写成一种字符串形式，非常强大。比如：
a = tf.random.normal([5, 3, 4]) b = tf.random.normal([5, 4, 6]) # 批量矩阵乘法 c1 = tf.einsum('bij,bjk->bik', a, b) # 转置 + 矩阵乘法 c2 = tf.einsum('bij->bji', a) # 更高级：外积展开 c3 = tf.einsum('bi,bj->bij', a[0], b[0])我第一次用einsum时觉得这是炫技，后来真香了。因为它把复杂的维度排列规则压缩到一行字符串里，不用仔细思考维度变换，写起来不容易错，而且TensorFlow对einsum做了算子调度优化。在实现注意力机制里的Q @ K^T、特征向量的高阶交互时，einsum几乎是标准写法。
如果矩阵乘法的带宽足够，einsum不会比手动matmul快太多，但它的优势是表达力。我在项目里一个重要的实践原则是：复杂维度变换尽量用einsum表达，简单线性层用tf.matmul，逐元素操作用*。三种操作定位清晰，代码可读性会好很多。
4. 静态图与XLA：从小步慢跑到全速冲刺
4.1 eager 与 graph：两种执行模式的效率鸿沟
TensorFlow 2.x默认是eager模式（动态图），代码写起来很舒服，一行一行立刻执行，调试方便。但代价是性能差。为什么？因为eager模式下，每个算子独立调度，框架要反复进出生效上下文、分配临时内存、等待CPU/GPU同步。对小型算子而言，调度开销甚至超过计算本身。
graph模式（静态图）则不同。它会先把整个计算流程构建成一张图，然后整图执行。图里的算子可以融合、复用内存、批量调度，尤其适合GPU这种大规模并行设备。
# eager 模式，逐算子执行 def eager_func(x, y, w): return tf.nn.relu(tf.matmul(x, w) + y) # graph 模式，整体编译执行 @tf.function def graph_func(x, y, w): return tf.nn.relu(tf.matmul(x, w) + y)从eager切到graph，对小型函数往往有立竿见影的效果。我用一个非常简单的矩阵乘法加激活测试过，同样的输入，graph模式比eager模式快一到两个数量级。原因很简单：算子调度开销被摊薄了，而且图模式能利用内核融合技术把多次数据读写合并成一次。
但要注意，tf.function不是银弹。它会在首次调用时对输入签名做trace（构建图），第二次调用相同签名时直接复用。如果输入签名变化太频繁，比如每次传入的batch大小都不同，函数会反复重新trace，反而更慢。所以果断提供input_signature，强制限定输入张量的shape和dtype，这是稳定性能的关键一步。
4.2 使用 tf.function 的三个隐藏规则
我在工程实践中总结出tf.function有三个容易踩的坑：
第一个坑：不要在tf.function内部创建tf.Variable。变量是全局状态，图执行时会为每次trace创建不同变量，结果就和预期完全不一样。正确做法是在函数外面创建变量，函数里只读取和更新。
第二个坑：普通的Pythonprint在tf.function内部不会在每次执行时都打印，它只会在trace阶段被记录下来。想真正跟踪运行时值，使用tf.print。但要注意，大量tf.print会强制图在运行过程中同步，性能影响很大。平时调试可以用，正式训练时务必关掉。
第三个坑：函数内部尽量不要依赖Python的全局变量。tf.function会捕获全局变换时的值作为常量，如果你后续改了全局值，图里用的还是旧值。这会导致你改了参数但训练行为没变化，排查起来非常痛苦。应该把这种需要动态变化的值作为参数显式传入函数。
@tf.function(input_signature=[ tf.TensorSpec(shape=[None, 10], dtype=tf.float32), tf.TensorSpec(shape=[], dtype=tf.float32) ]) def process(x, alpha): return tf.nn.relu(x) * alpha这里的alpha是标量张量，每次传入新值，图都能根据实际值计算。
4.3 开启XLA：让算子自动融合
XLA（Accelerated Linear Algebra）是TensorFlow自带的高性能编译器。它的核心能力是算子融合：把相邻的多个算子合并成单个融合算子，减少中间结果的写回和读入。典型的例子是把matmul + add + relu融合成一个算子，一次执行完成。这在GPU上效果极佳，因为减少了显存读写次数，也减少了kernel启动次数。
开启方法最简单的是在代码开头：
tf.config.optimizer.set_jit(True)也可以通过环境变量TF_XLA_FLAGS=--tf_xla_auto_jit=2实现，二档会自动尝试对所有可用算子做编译。
我有一次在某个图像处理Demo中开启XLA后，整体耗时下降了近四成。但也要泼一盆冷水：XLA不是所有操作都加速，它对计算密集、形状固定的大算子效果显著，对大量小算子、动态shape、控制流复杂的代码反而可能增加编译开销。一般在模型训练或推理的主干网络层效果明显，而在数据处理阶段用了反而得不偿失。
XLA还有一个隐藏优势：它能做形状推断和内存优化，让显存分配更紧凑。在显存资源紧张时，开启XLA有时能减少out of memory的发生概率。但也不是保证，该爆显存还是会爆。
当使用动态shape时，XLA可能需要额外的重新编译，跑一段时间后你会发现一个奇怪的停顿，那就是它在编译新的形态组合。想避免的话，尽量让输入shape固定，或者用tf.function的input_signature固定签名。
5. 高性能张量流水线的实战打磨
5.1 控制显存分配和设备放置
高性能计算的前提，是搞清楚张量到底在哪个设备上。很多训练慢的场景，问题不在模型，而在数据反复在CPU和GPU之间搬移。
一个经典的错误是：在循环的每一轮里把数据用.numpy()转回Python，做点预处理，再转回Tensor。这一转是纯CPU操作，而且会阻塞GPU流水线。正确做法是所有预处理都在TensorFlow计算图内完成。
查看张量所在设备，用print(x.device)。显存分配方面，TensorFlow默认按需增长，但有些环境会一次性把所有显存占满。想要按需增长可以设置：
gpus = tf.config.experimental.list_physical_devices('GPU') if gpus: for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True)还有一个实战中很重要但常被忽略的点：设set_log_device_placement(True)可以打印每个算子的设备分配，用来检查是否有意外的CPU算子混入GPU计算流。
tf.debugging.set_log_device_placement(True)如果一个模型被拆得非常碎，包含大量小算子，那么每次kernel启动的固定开销会成为瓶颈。这种情况下，可以尝试把多个小算子合并成一个自定义大算子，或者干脆增大单次计算的数据量，用更大batch来提升每个算子的数据规模，摊薄启动开销。
5.2 用 tf.data 构建高效数据流水线
张量操作的性能瓶颈往往不在计算，而在数据喂入。tf.data.Dataset是官方推荐的数据管线方案，它有两个机制对性能提升至关重要：预取（prefetch）和并行映射（map）。
prefetch的原理是数据读取和模型计算重叠。GPU在算当前batch时，CPU已经在准备下一个batch的数据，两者并行，消除等待。
并行map同理：map函数可以对多个样本同时做预处理，而不是逐样本串行。
dataset = tf.data.Dataset.from_tensor_slices((features, labels)) dataset = dataset.shuffle(10000) dataset = dataset.batch(128) dataset = dataset.map(preprocess, num_parallel_calls=tf.data.AUTOTUNE) dataset = dataset.prefetch(tf.data.AUTOTUNE)这里的顺序是有讲究的：先shuffle再batch，否则shuffle粒度会变成batch级别，打乱效果大打折扣。batch之后再做map，虽然处理的是整个batch，但适用场景较少，一般还是先map后batch，因为逐样本预处理通常更高效。
我在实践中的一个体会：AUTOTUNE让框架自动调整并行线程数，比自己硬编码数值效果更好。但有两个例外：如果数据集已经被完全打乱并装在内存里，num_parallel_calls过大反而引起CPU争抢；如果预处理本身极简单（只是归一化），并行开销可能比串行还高。这种时候需要手动压测，找到一个合适的并发度。
5.3 混合精度与算子布局优化
混合精度是近年来最直接的性能红利之一。核心思路是用精度较低的浮点数格式参与部分运算，在几乎不影响模型精度的情况下，大幅提升计算速度并减少内存。
tf.keras.mixed_precision提供了便捷的全局策略设置：
from tensorflow.keras import mixed_precision policy = mixed_precision.Policy('mixed_float16') mixed_precision.set_global_policy(policy)设置后，变量仍然以float32存储（保证数值稳定），但实际计算用float16或bfloat16，梯度更新前自动转回float32。这带来的效果有两层：一是float16的数据只有float32一半大小，同样的显存能装下更大的batch；二是建立在GPU上，float16的算力往往远高于float32，计算本身更快。
这里有一个我踩过的坑：混合精度下，某些算子比如softmax、LayerNorm必须要用float32计算，否则数值会溢出。官方实现会在内部自动处理，但如果你手写自定义层，就要自己加上tf.cast(x, tf.float32)再算，最后转回float16。我最初做自注意力模块时没注意这个细节，模型训练几天后loss变成NaN，查了很久才定位到是half精度下注意力权重的softmax数值溢出。
另一个优化点是张量布局。CPU友好的是NHWC（通道在后），GPU常用的是NCHW。TensorFlow各层自动适配，但如果你手写低层特征处理，要特别注意布局。不要每层都换来换去，我在项目中的原则是输入层明确一个统一布局，后续所有自定义操作维持不动。
6. 一个真实案例：特征提取流程从1.2秒降到0.3秒
作为综合演示，我构建了一个典型的离线特征提取流程：从本地读取一大批图像，经过预处理、模型推理、特征后处理，最终输出固定维度的特征向量。这是很多检索系统、推荐系统的标准前置环节。我在测试环境里输入是指定的图片数据集，代码是模拟的，但优化思路完全通用。
第一版是纯eager写法，直接循环处理每一张图：
def process_images_eager(images): results = [] for img in images: # 预处理 x = tf.cast(img, tf.float32) / 255.0 x = tf.image.resize(x, [224, 224]) x = tf.expand_dims(x, axis=0) # 推理 feat = model(x) # 假设模型输出 [1, 512] results.append(feat[0]) return tf.stack(results)这个版本的瓶颈很典型：每一张图都要单独经过模型推理，模型内部算子的启动开销被放大到极致。实测1.2秒处理完这批图片。
第二版用tf.function包住整个处理循环后，耗时降到0.8秒左右。原因是算子调度开销大幅缩小，但这还不够，因为数据处理和GPU计算仍然是串行。
第三版引入tf.data流水线，批处理加prefetch：
dataset = tf.data.Dataset.from_tensor_slices(images) dataset = dataset.map(preprocess, num_parallel_calls=tf.data.AUTOTUNE) dataset = dataset.batch(32) dataset = dataset.prefetch(tf.data.AUTOTUNE) @tf.function(input_signature=[ tf.TensorSpec(shape=[None, 224, 224, 3], dtype=tf.float32) ]) def infer_batch(batch): return model(batch) for batch in dataset: feats = infer_batch(batch) all_feats.append(feats)耗时降到0.6秒。提升来自三个层面：预处理与模型推理重叠、模型总计算从逐张变成批量、批量内算子比单张嵌入并行度更高。
第四版叠加混合精度，耗时降到0.42秒左右。特征提取这类推理任务，半精度造成的精度损失在我的场景里几乎没有影响，但速度提升非常明显。
第五版再开XLA，耗时继续降到0.31秒。到了这个阶段，继续优化的边际收益非常小，再往下就需要在模型结构层面做改动，或者采用量化部署。
我把这组优化记录在一个小表格里，方便自己后续复盘：
| 优化阶段 | 耗时（秒） | 主要改动 | 提速原因 |
|---|---|---|---|
| 基线eager循环 | 1.2 | 逐张图推理 | 算子启动开销大，数据串行 |
| 加入tf.function | 0.8 | 整流程静态图化 | 缩减调度开销 |
| 构建数据流水线 | 0.6 | Dataset批量+prefetch | 数据加载与计算重叠 |
| 开启混合精度 | 0.42 | 半精度计算 | 降低带宽与计算量 |
| 开启XLA | 0.31 | 算子融合编译 | 减少中间结果读写 |
这个案例的普适意义在于：优化不是某一个魔法开关，而是多级组合。每一步都只解决一个明确的瓶颈，叠加起来才有质的飞跃。我的实际经验是，先做profiling，找到瓶颈在IO还是在计算，然后再针对性优化。
7. 常见错误速查表：避开我踩过的那些张量坑
最后把我在实际使用中遇到的高频问题整理成一张速查表，希望能帮你快速定位问题。
| 报错信息 | 常见原因 | 解决方案 |
|---|---|---|
| Incompatible shapes for broadcasting | 两个张量维度不兼容 | 检查算式两边的shape，用expand_dims或reshape对齐维度 |
| Shape must be rank N but is rank M | 传入算子的张量维度和预期不符 | 打印x.shape确认，必要时增加squeeze/expand_dims |
| Resource exhausted: OOM when allocating tensor | 显存或内存不足 | 减小batch、开启显存按需增长、考虑混合精度 |
| Variable not found in graph | tf.function内部创建了变量 | 把变量移到函数外部创建 |
| Check failed: dtype mismatch | 两个张量类型不一致 | 用tf.cast统一dtype，常用的是转float32 |
| Failed to copy tensor to device | CPU张量和GPU张量混合运算 | 检查device，显式用tf.identity或tf.device控制设备 |
除了这些常见报错，还有两个我吃过亏的方向值得特意提一下。
一个是tf.TensorShape的判等问题。x.shape == [None, 224, 224, 3]这种写法在动态shape时可能得到False，因为None不等于任何具体整数。更稳的判等是用x.shape.as_list() == [None, 224, 224, 3]，或者直接用x.shape.is_compatible_with。
另一个是tf.function内部的Python控制流。如果你在函数里写了if语句，并且条件依赖张量值，TensorFlow会把if转换成tf.cond，这在大循环里效率不错，但条件分支一旦很多，图会变得非常庞大。此时可以考虑把复杂控制流拆出来，在Python层做，只在需要高性能的核心计算段用tf.function。
我自己后期的一个习惯是：先写纯Python零食代码验证逻辑，再用tf.function包装最内层热路径，最后用input_signature固定签名。这样既保住开发效率，又拿到运行性能。
做张量操作优化这件事，最关键的往往不是记住某个具体函数，而是建立一种直觉：每个操作背后，数据以什么形式、在哪个设备上流动，开销花在了哪里。带着这个认知去排查性能问题，你很快会发现，绝大多数瓶颈都能被归因到设备拷贝、算子调度、布局转换这三类问题上。搞定了这三类问题，不管你的模型是简易Demo还是生产系统，性能都会有肉眼可见的提升。
