TensorFlow变长数据Batch处理:TFRecordDataset与padded_batch实战

竹里
竹里
发布时间:2026-10-13| 6人浏览过
原创

本文详解TensorFlow中如何处理长度不一的变长数据(如语音、视频或NLP样本)。通过对比直接padding写入与使用padded_batch方法,重点解析padded_batch的参数配置及tf.VarLenFeature配合sparse_tensor_to_dense的使用,结合MNIST示例展示完整代码实现。

变长数据Batch读取的挑战与方案

在前文《TensorFlow入门:TFRecord与TFRecordDataset的使用》中,我们介绍了利用dataset的batch方法对TFRecord文件进行批量读取。然而,当数据集中每条记录的维度不一致时(这在语音识别、视频处理及自然语言处理领域十分常见),直接调用batch方法将无法正常工作。针对这一场景,通常有两种解决方案:

第一种方案是在将数据写入TFRecord之前,先对所有数据进行填充(padding)至统一长度,然后再写入文件。这种做法的主要缺陷在于,如果大量数据的实际长度远小于最大长度,将会导致存储空间的极大浪费。

第二种方案是利用dataset提供的padded_batch方法。该方法通过参数padded_shapes指定每条记录中各个成员需填充到的形状:若成员为标量(scalar),使用[]表示;若为列表(list),使用[mx_length]表示;若为数组(array),则使用[d1,...,dn]表示。例如,当各成员顺序依次为标量、列表和数组时,padded_shapes应设置为([], [mx_length], [d1,...,dn])。该函数的具体说明如下:

padded_batch方法详解


padded_batch(
 batch_size,
 padded_shapes,
 padding_values=None #默认使用各类型数据的默认值,一般使用时可忽略该项
)

为了直观说明,我们使用MNIST数据集进行演示。首先,在将MNIST数据写入TFRecord之前,我们需要对数据进行预处理,使每个MNIST图像的尺寸变得不相等,具体操作如下:

生成变长TFRecord数据


import tensorflow as tf
from tensorflow.contrib.learn.python.learn.datasets.mnist import read_data_sets
 
mnist = read_data_sets("MNIST_data/", one_hot=True)
 
 
def get_tfrecords_example(feature, label):
 tfrecords_features = {}
 feat_shape = feature.shape
 tfrecords_features['feature'] = tf.train.Feature(float_list=tf.train.FloatList(value=feature))
 tfrecords_features['shape'] = tf.train.Feature(int64_list=tf.train.Int64List(value=list(feat_shape)))
 tfrecords_features['label'] = tf.train.Feature(float_list=tf.train.FloatList(value=label))
 return tf.train.Example(features=tf.train.Features(feature=tfrecords_features))
 
 
def make_tfrecord(data, outf_nm='mnist-train'):
 feats, labels = data
 outf_nm += '.tfrecord'
 tfrecord_wrt = tf.python_io.TFRecordWriter(outf_nm)
 ndatas = len(labels)
 print(feats[0].dtype, feats[0].shape, ndatas)
 assert len(labels[0]) > 1
 for inx in range(ndatas):
 ed = random.randint(0,3) #随机丢掉几个数据点,以使长度不等
 exmp = get_tfrecords_example(feats[inx][:-ed], labels[inx])
 exmp_serial = exmp.SerializeToString()
 tfrecord_wrt.write(exmp_serial)
 tfrecord_wrt.close()
 
import random
nDatas = len(mnist.train.labels)
inx_lst = range(nDatas)
random.shuffle(inx_lst)
random.shuffle(inx_lst)
ntrains = int(0.85*nDatas)
 
# make training set
data = ([mnist.train.images[i] for i in inx_lst[:ntrains]], \
 [mnist.train.labels[i] for i in inx_lst[:ntrains]])
make_tfrecord(data, outf_nm='mnist-train')
 
# make validation set
data = ([mnist.train.images[i] for i in inx_lst[ntrains:]], \
 [mnist.train.labels[i] for i in inx_lst[ntrains:]])
make_tfrecord(data, outf_nm='mnist-val')
 
# make test set
data = (mnist.test.images, mnist.test.labels)
make_tfrecord(data, outf_nm='mnist-test')

在使用dataset加载批量数据时,解析过程需采用tf.VarLenFeature(tf.datatype)替代tf.FixedLenFeature([], tf.datatype),并必须配合tf.sparse_tensor_to_dense函数使用,具体代码如下:

解析与加载数据集


import tensorflow as tf
 
train_f, val_f, test_f = ['mnist-%s.tfrecord'%i for i in ['train', 'val', 'test']]
 
def parse_exmp(serial_exmp):
 feats = tf.parse_single_example(serial_exmp, features={'feature':tf.VarLenFeature(tf.float32),\
 'label':tf.FixedLenFeature([10],tf.float32), 'shape':tf.FixedLenFeature([], tf.int64)})
 image = tf.sparse_tensor_to_dense(feats['feature']) #使用VarLenFeature读入的是一个sparse_tensor,用该函数进行转换
 label = tf.reshape(feats['label'],[2,5]) #把label变成[2,5],以说明array数据如何padding
 shape = tf.cast(feats['shape'], tf.int32)
 return image, label, shape
 
def get_dataset(fname):
 dataset = tf.data.TFRecordDataset(fname)
 return dataset.map(parse_exmp) # use padded_batch method if padding needed
 
epochs = 16
batch_size = 50 
padded_shapes = ([784],[3,5],[]) #把image pad至784,把label pad至[3,5],shape是一个scalar,不输入数字
# training dataset
dataset_train = get_dataset(train_f)
dataset_train = dataset_train.repeat(epochs).shuffle(1000).padded_batch(batch_size, padded_shapes=padded_shapes)

以上关于TensorFlow中TFRecordDataset处理变长数据Batch读取的详解内容分享完毕,希望能为读者提供有益参考,也欢迎大家继续支持本站。

相关标签:

免责声明:本文内容仅供参考,不代表本站观点。转载或引用请注明出处。

Python 相关产品和应用

全部下载
Selenium

Selenium

4.51.0 · Windows

用于浏览器自动化和 Web 测试的 Python 工具。

scikit-learn

scikit-learn

1.9.1 · Windows

常用 Python 机器学习算法与模型评估工具库。

Black

Black

26.10.1 · Windows

强调一致性的 Python 自动代码格式化工具。

Ruff

Ruff

0.17.1 · macOS

由 Rust 编写的高速 Python 代码检查和格式化工具。

pytest

pytest

9.1.1 · Windows

简洁而强大的 Python 自动化测试框架。

FastAPI

FastAPI

0.146.0 · Windows

高性能 Python API 框架,支持类型提示和自动接口文档。

Flask

Flask

3.1.3 · Windows

轻量灵活的 Python Web 应用框架。

Django

Django

6.1.2 · Windows

成熟的 Python Web 框架,适合构建内容站点和业务系统。

Spyder

Spyder

6.1.7 · Windows

面向科学计算与数据分析的 Python 集成开发环境。

SciPy

SciPy

1.18.1 · Windows

面向科学、工程和技术计算的 Python 算法库。