import numpy as np import tensorflow as tf import sys from preprocess_data import get_dataset_1M import os from pathlib import Path p = Path(__file__).parents[1] ROOT_DIR=os.path.abspath(os.path.join(p, '..', 'data/')) TF_RECORD_TRAIN_PATH='/tf_records/train' TF_RECORD_TEST_PATH='/tf_records/test' def _add_to_tfrecord(data_sample,tfrecord_writer): data_sample=list(data_sample.astype(dtype=np.float32)) example = tf.train.Example(features=tf.train.Features(feature={'movie_ratings': float_feature(data_sample)})) tfrecord_writer.write(example.SerializeToString()) def _get_output_filename(output_dir, idx, name): return '%s/%s_%03d.tfrecord' % (ROOT_DIR+output_dir, name, idx) def int64_feature(value): if not isinstance(value, list): value = [value] return tf.train.Feature(int64_list=tf.train.Int64List(value=value)) def float_feature(value): if not isinstance(value, list): value = [value] return tf.train.Feature(float_list=tf.train.FloatList(value=value)) def bytes_feature(value): if not isinstance(value, list): value = [value] return tf.train.Feature(bytes_list=tf.train.BytesList(value=value)) def main(): ''' Writes the .txt training and testing data into binary TF_Records.''' SAMPLES_PER_FILES=100 training_set, test_set=get_dataset_1M() for data_set, name, dir_ in zip([training_set, test_set], ['train', 'test'], [TF_RECORD_TRAIN_PATH, TF_RECORD_TEST_PATH]): num_samples=len(data_set) i = 0 fidx = 1 while i < num_samples: tf_filename = _get_output_filename(dir_, fidx, name=name) with tf.python_io.TFRecordWriter(tf_filename) as tfrecord_writer: j = 0 while i < num_samples and j < SAMPLES_PER_FILES: sys.stdout.write('\r>> Converting sample %d/%d' % (i+1, num_samples)) sys.stdout.flush() sample = data_set[i] _add_to_tfrecord(sample, tfrecord_writer) i += 1 j += 1 fidx += 1 print('\nFinished converting the dataset!') if __name__ == "__main__": main()
这段代码是用来将电影评分数据集转换为 TensorFlow 训练所需的二进制 TFRecord 格式的。这里采用的是 MovieLens 数据集，其中包含了 1 百万个电影评分记录，用于推荐系统任务的训练和测试。
该代码主要分为几个部分：
- 引入必要的 Python 模块和工具
- 定义一些用于将数据样本添加到 TFRecord 文件中的函数，以及获取输出文件名的函数
- 定义一些特征转换函数，如 int64_feature、float_feature 和 bytes_feature，这些函数用于将样本的各个特征值转换为 TFRecord 文件所需的数据类型
- 定义主函数 main，其中首先调用 get_dataset_1M 函数，获取训练和测试集数据。然后，针对训练集和测试集，分别循环将数据集写入 TFRecord 文件中。每个文件中包含 100 个数据样本，这是为了方便后续处理和管理。在将数据样本添加到文件中时，调用了 _add_to_tfrecord 函数，并传入当前样本和 TFRecordWriter 对象，以将样本写入文件中。在完成一个文件的写入后，使用 fidx += 1 更新文件编号，以准备开始下一个文件的写入。最后输出完成的信息。
- 在主函数末尾，通过 if name == "main" 判断当前文件是否为主文件，如果是则执行 main 函数。
这段代码的主要目的是将原始数据转换为可以用于训练机器学习模型的格式，即 TFRecord 格式。这种格式不仅可以更快地读取和处理数据，而且可以节省存储空间。
这段代码中，main() 函数是主要的执行函数，它将训练和测试数据集写入到二进制的 TF_Records 文件中。在 main() 函数中，首先通过 get_dataset_1M() 函数获取训练和测试数据集。接下来，对于每个数据集，程序会将数据分成若干份，每份的大小为 SAMPLES_PER_FILES=100。对于每一份数据，程序会创建一个对应的 TF_Records 文件，将数据写入该文件中。
在写入数据时，程序使用 _add_to_tfrecord() 函数将每个数据样本转换为一个 TF_Records 格式，并将其写入 TF_Records 文件中。其中，_add_to_tfrecord() 函数将每个数据样本转换为一个 tf.train.Example 对象，该对象包含一个 tf.train.Features 对象，其中包含一个 movie_ratings 特征。movie_ratings 特征的值是一个由浮点数构成的列表，列表中的每个元素对应一个电影的评分。
在转换数据样本时，程序使用 float_feature() 函数将浮点数列表转换为一个 tf.train.Feature 对象，该对象将浮点数列表转换为一个 tf.train.FloatList 对象。在写入数据时，程序使用 tf.python_io.TFRecordWriter() 函数创建一个 TF_Records 文件写入器，并将数据写入该文件中。最后，程序输出转换数据集的信息。
