如何高效从多份json.gz文件创建tf.dataset并解决类型报错？

错误原因分析 #

你遇到的InternalError: Unsupported object type dict核心原因有两个：

1. tf.numpy_function指定的返回类型(tf.int64)和实际返回的字典结构完全不匹配，TensorFlow无法解析这种类型冲突的返回值。
2. 直接返回pandas生成的字典，无法被tf.data.Dataset.from_tensor_slices正确转换为张量数据集。

修正后的完整代码 #

import tensorflow as tf
import pandas as pd
import gzip
import json
from datetime import datetime

# 替换为你实际的主题映射字典
themes = {"computer science": 0, "physics": 1, "biology": 2}

def load_file(file_path):
    # 将TensorFlow字符串张量转为Python可处理的字符串
    file_path = file_path.numpy().decode('utf-8')
    
    with gzip.open(file_path, 'r') as fin:
        json_bytes = fin.read()
    json_str = json_bytes.decode('utf-8')
    data = json.loads(json_str)
    
    # 规范化JSON数据，修正路径参数格式
    bb = pd.json_normalize(
        data,
        record_path='items',
        meta=['indexed', ['title', 'publisher', 'type', 'indexed.date-parts', 'subject']],
        errors='ignore'
    )
    
    # 过滤缺失关键字段的样本
    bb.dropna(subset=['title', 'publisher', 'type', 'indexed.date-parts', 'subject'], inplace=True)
    
    # 主题标签映射逻辑优化
    def map_subject(subject_list):
        intersect = set(subject_list) & set(themes.keys())
        return themes[next(iter(intersect))] if intersect else len(themes) + 1
    bb['subject'] = bb['subject'].apply(map_subject)
    
    # 处理标题和日期格式
    bb['title'] = bb['title'].str.join('')
    bb['indexed.date-parts'] = bb['indexed.date-parts'].apply(
        lambda tpl: datetime.strptime('-'.join(map(str, tpl[0])), '%Y-%m-%d').strftime('%Y-%m-%d')
    )
    
    # 转换为TensorFlow可识别的numpy数组结构
    return {
        'title': bb['title'].values.astype('U'),
        'publisher': bb['publisher'].values.astype('U'),
        'type': bb['type'].values.astype('U'),
        'indexed_date': bb['indexed.date-parts'].values.astype('U'),
        'subject': bb['subject'].values.astype('int64')
    }

def tf_load_file(file_path):
    # 用tf.py_function包装Python函数，明确指定每个字段的张量类型
    return tf.py_function(
        func=load_file,
        inp=[file_path],
        Tout={
            'title': tf.string,
            'publisher': tf.string,
            'type': tf.string,
            'indexed_date': tf.string,
            'subject': tf.int64
        }
    )

# 文件列表
file_list = [
    'file_2021_01/10625.json.gz',
    'file_2021_01/23897.json.gz',
    'file_2021_01/12169.json.gz',
    'file_2021_01/427.json.gz'
]

# 构建文件数据集并打乱顺序
filenames = tf.data.Dataset.from_tensor_slices(file_list)
filenames = filenames.shuffle(buffer_size=len(file_list))

# 并行加载文件并实现混洗
dataset = filenames.interleave(
    map_func=lambda x: tf.data.Dataset.from_tensor_slices(tf_load_file(x)),
    cycle_length=tf.data.AUTOTUNE,
    num_parallel_calls=tf.data.AUTOTUNE,
    deterministic=False  # 关闭确定性加载提升效率
)

# 按特定标签过滤数据（示例：只保留主题为0的样本）
target_subject = 0
dataset = dataset.filter(lambda x: tf.equal(x['subject'], target_subject))

# 全局样本混洗（根据内存情况调整buffer_size）
dataset = dataset.shuffle(buffer_size=10000)

# 批量处理+预取优化
dataset = dataset.batch(32)
dataset = dataset.prefetch(tf.data.AUTOTUNE)

关键修正与优化说明 #

• 类型匹配修复：用tf.py_function替代tf.numpy_function，明确声明每个返回字段的张量类型，解决字典类型不兼容问题。
• 数据格式转换：将pandas数据转为指定类型的numpy数组，确保TensorFlow能正确解析为张量。
• 并行效率提升：使用tf.data.AUTOTUNE让框架自动调整并行度，deterministic=False进一步优化加载速度。
• 过滤与混洗实现：通过filter方法完成主题过滤，先做文件级混洗再做样本级混洗，保证数据分布均匀。
• 依赖与参数修正：补充了缺失的gzip、json、datetime导入，修正了pd.json_normalize的参数格式，避免解析错误。

额外优化建议 #

• 若文件数量极大，shuffle(buffer_size=len(file_list))可能占用过多内存，可改用分阶段混洗策略。
• 针对超大单个文件，可在load_file中实现分块读取，避免内存溢出。
• 可将主题映射逻辑转为纯TensorFlow操作，减少Python函数调用开销。

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