如何将pandas DataFrame转换为TensorFlow数据集用于文本二分类

安全阿岳

2026-10-1

完全可以实现，你可以用tf.data.Dataset.from_tensor_slices方法完成转换，最终得到的数据集格式和官方教程中的输出完全兼容，可直接对接后续BERT训练流程，步骤如下：

步骤1：拆分数据集

首先按照官方教程的拆分比例，将完整DataFrame拆分为训练集、验证集、测试集：

import tensorflow as tf
import pandas as pd
from sklearn.model_selection import train_test_split

# 此处替换为你已经加载好的DataFrame变量名
df = pd.read_csv("你的数据集路径.csv")

AUTOTUNE = tf.data.AUTOTUNE
batch_size = 32
seed = 42

# 先拆分出20%作为测试集
train_val_df, test_df = train_test_split(df, test_size=0.2, random_state=seed)
# 再从剩余数据中拆分20%作为验证集，和官方教程划分逻辑一致
train_df, val_df = train_test_split(train_val_df, test_size=0.2, random_state=seed)

步骤2：转换为TensorFlow数据集

直接通过from_tensor_slices方法生成和官方格式一致的数据集：

# 生成训练集：(文本批次, 标签批次) 结构
train_ds = tf.data.Dataset.from_tensor_slices(
    (train_df['Description'].values, train_df['Target'].values)
).shuffle(buffer_size=len(train_df), seed=seed).batch(batch_size).cache().prefetch(buffer_size=AUTOTUNE)

# 生成验证集
val_ds = tf.data.Dataset.from_tensor_slices(
    (val_df['Description'].values, val_df['Target'].values)
).batch(batch_size).cache().prefetch(buffer_size=AUTOTUNE)

# 生成测试集
test_ds = tf.data.Dataset.from_tensor_slices(
    (test_df['Description'].values, test_df['Target'].values)
).batch(batch_size).cache().prefetch(buffer_size=AUTOTUNE)

# 手动定义分类名称，功能和官方教程输出的class_names完全一致
class_names = ['failed', 'passed']