TensorFlow 2.9中image.resize()等操作耗时过高的排查与优化请求

我正在用TensorRT模型在GPU上做图像分割预测，输出对应图像的分割掩码。用argmax和image.resize把掩码恢复成原始图像尺寸时，发现其他TensorFlow函数单次耗时都低于3ms，但tf.image.resize()耗时高达50ms，达不到项目要求的掩码处理总耗时约20ms的标准。

代码示例 #

prd_tf = model(img_tf)

# Segmentation output mask
msk_tf = tf.argmax(prd_tf['conv2d_19'], axis=-1) # < 3ms                            
msk_tf = tf.reshape(msk_tf, (512, 512, 1)) # < 3ms                                     
msk_tf = tf.image.resize(msk_tf, (img_height, img_width), method='nearest') # > 40 ms
msk_tf = tf.reshape(msk_tf, (img_height, img_width)) # < 3ms                          

我尝试过两种优化方案，但效果都不理想：

• 修改数据类型为int16：argmax耗时超过40ms，resize等操作耗时低于3ms；
• Numpy替代方案：调用.numpy()操作耗时超40ms，其他操作耗时正常。

问题 #

1. 为什么会出现这种耗时异常的情况？
2. 有哪些方案能把输出掩码的处理耗时控制在20ms以内？

问题解答 #

1. 耗时异常的原因分析 #

• TensorFlow的resize对整数张量优化不足：tf.argmax输出的是int32类型张量，而tf.image.resize针对整数类型的GPU kernel优化远不如浮点型，此时可能触发CPU fallback或低效GPU计算，导致耗时飙升；改成int16后argmax变慢，是因为TensorFlow对int16类型的argmax支持不完善，GPU kernel优化程度低。
• GPU-CPU数据传输开销：调用.numpy()会将GPU上的张量同步拷贝到CPU，跨设备数据传输的延迟远高于计算本身，直接导致耗时超40ms。
• TensorRT与TensorFlow交互的额外开销：如果是通过TF-TRT集成的TensorRT模型，张量从TRT输出到TF张量的过程中，数据布局或类型未做最优适配，导致后续TF操作无法高效利用GPU算力。

2. 可行的优化方案 #

方案1：把后处理逻辑集成到TensorRT模型内

将argmax和resize操作直接加入TensorRT推理图，全程在GPU上完成，避免跨框架开销：

• 导出模型时，将argmax和resize逻辑加入计算图，再转换为TensorRT引擎；
• 使用TensorRT Python API手动构建包含ArgMax和Resize插件的推理管线，这两个插件对整数类型的支持效率远高于TensorFlow。

方案2：优化TensorFlow内的resize操作

• 先转浮点型再resize：将argmax输出的int32张量转成float32，完成resize后再转回int32。TensorFlow对浮点型的resize GPU kernel优化成熟，耗时可降至3ms以内。示例代码：

msk_tf = tf.argmax(prd_tf['conv2d_19'], axis=-1) # <3ms
msk_tf = tf.cast(msk_tf, tf.float32)
msk_tf = tf.expand_dims(msk_tf, axis=-1)
msk_tf = tf.image.resize(msk_tf, (img_height, img_width), method='nearest') # ~2ms
msk_tf = tf.cast(tf.squeeze(msk_tf, axis=-1), tf.int32) # <3ms

• 直接调用底层原始操作：用tf.raw_ops.ResizeNearestNeighbor替代高层API，减少额外的参数检查和转换开销，进一步降低耗时。

方案3：用CuPy替代Numpy做GPU端处理

如果必须做CPU端处理，改用CuPy直接在GPU上执行数组操作，避免跨设备数据拷贝：

import cupy as cp

msk_cp = cp.asarray(msk_tf) # TF与CuPy零拷贝桥接
msk_cp = cp.resize(msk_cp, (img_height, img_width)) # GPU上执行，耗时<3ms
msk_tf = tf.convert_to_tensor(msk_cp)

方案4：调整模型输出分辨率

如果业务允许，修改模型的输出分辨率为原始图像尺寸，直接省去resize步骤，这是最彻底的优化方式。

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