tf.data Dataset批处理后映射预处理函数报错的解决方法

问题描述 #

在图像分类项目中尝试遵循TensorFlow官方“向量化映射”建议，将预处理放在tf.data.Dataset批处理之后执行，代码结构如下：

ds = ds.batch(batch_size)
ds = ds.map(process_batch)

但运行时触发InaccessibleTensorError，报错提到循环内张量超出作用域不可访问，且代码中并未显式写while循环。

环境配置与基础代码：

import numpy as np
import os
import tensorflow as tf
import glob

data_path = "data/img_data"
img_height = 64
img_width = 64
AUTOTUNE = tf.data.AUTOTUNE
batch_size=32

img_files = glob.glob(f"{data_path}/*/*.jpg")
n_imgs = len(img_files)

ds = tf.data.Dataset.list_files(img_files,shuffle=False)
ds = ds.shuffle(n_imgs,reshuffle_each_iteration=False)

class_names = [i.split("/")[-1] for i in glob.glob(f"{data_path}/*")]

单样本预处理可正常运行，但批处理版process_batch报错：

@tf.function
def process_batch(batch):
    batch_labels = []
    batch_imgs = []
    for i in batch:
        
        label = tf.strings.split(i,os.sep)[-2]
        label = tf.argmax(label==class_names)
        batch_labels.append(label)
        
        img = tf.io.read_file(i)
        img = tf.io.decode_jpeg(img, channels=3)
        img = tf.image.resize(img,[img_height,img_width])
        img = tf.cast(img,tf.float32)/255
        batch_imgs.append(img)        
    
    batch_labels = tf.convert_to_tensor(batch_labels, dtype=tf.int64)
    batch_imgs = tf.convert_to_tensor(batch_imgs, dtype=tf.float32)
    return imgs,labels  # 此处变量名错误

def config_ds2(ds):
    ds = ds.shuffle(buffer_size=ds.cardinality().numpy())
    ds = ds.batch(batch_size,drop_remainder=True)
    ds = ds.map(process_batch)
    return ds

ds2 = config_ds2(ds)

报错信息：

InaccessibleTensorError: in user code:

    File "/var/folders/2c/cr8tgk091dg1qcqlnkxypj5m0000gn/T/ipykernel_17058/3752611221.py", 
line 54, in process_batch  *
        batch_labels = tf.convert_to_tensor(batch_labels, dtype=tf.int64)

    InaccessibleTensorError: <tf.Tensor 'while/ArgMax:0' shape=() dtype=int64> is out of scope and 
cannot be used here. Use return values, explicit Python locals or TensorFlow collections to access it.
    Please see 
https://www.tensorflow.org/guide/function#all_outputs_of_a_tffunction_must_be_return_values 
for more information.

The tensor <tf.Tensor 'while/ArgMax:0' shape=() dtype=int64> cannot be accessed from 
FuncGraph(name=process_batch, id=4966476864), because it was defined in 
FuncGraph(name=while_body_135, id=4967155888), which is out of scope.

错误原因分析 #

1. 隐式while循环：@tf.function会将Python的for i in batch转换为TensorFlow的while循环，循环内部创建的张量属于子图，外部无法直接访问，导致作用域错误。
2. 变量名错误：函数最后返回的imgs,labels未定义，正确应为batch_imgs,batch_labels。
3. 非向量化操作：用Python列表收集张量再转换的方式不符合TensorFlow向量化操作规范，既低效又容易引发作用域问题。

解决方案 #

使用TensorFlow原生的向量化API处理整个批次，完全避免Python循环，同时修正返回值错误：

@tf.function
def process_batch(batch):
    # 批量处理标签：分割文件路径获取类别名，转换为索引
    parts = tf.strings.split(batch, os.sep)
    class_names_tensor = tf.constant(class_names)
    labels = tf.argmax(tf.equal(parts[:, -2:][:, 0], class_names_tensor), axis=1)
    
    # 批量读取、解码、预处理图像
    imgs = tf.io.read_file(batch)
    imgs = tf.io.decode_jpeg(imgs, channels=3)
    imgs = tf.image.resize(imgs, [img_height, img_width])
    imgs = tf.cast(imgs, tf.float32) / 255.0
    
    return imgs, labels

def config_ds2(ds):
    ds = ds.shuffle(buffer_size=ds.cardinality().numpy())
    ds = ds.batch(batch_size, drop_remainder=True)
    ds = ds.map(process_batch, num_parallel_calls=AUTOTUNE)  # 增加并行调用提升效率
    return ds

# 验证数据集
ds2 = config_ds2(ds)
for batch_imgs, batch_labels in ds2.take(1):
    print(f"图像批次形状: {batch_imgs.shape}")
    print(f"标签批次形状: {batch_labels.shape}")

关键优化点 #

• 向量化标签处理：用tf.strings.split批量分割路径，结合tf.equal和tf.argmax批量获取标签索引，避免逐样本循环。
• 批量图像操作：tf.io.read_file、tf.io.decode_jpeg等API原生支持批量输入，无需手动循环处理每个文件。
• 并行调用：在map中加入num_parallel_calls=AUTOTUNE，让TensorFlow自动优化并行处理数量，提升流水线效率。

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