Spark下使用PySpark生成TFRecord样本文件并存储到HDFS

PySpark集成了TensorFlow的TFRecord格式数据的读取和写入支持。通过安装tfrecord2.12.3.0.jar包，我们可以在Spark集群中使用PySpark来生成TFRecord样本文件，并将其直接存储到HDFS上。以下是生成TFRecord样本文件的示例代码：

from pyspark.ml.feature import VectorAssembler

from pyspark.ml.linalg import Vectors

生成测试数据

data = [(Vectors.dense([1, 2, 3]), 1.0),

(Vectors.dense([4, 5, 6]), 2.0)]

df = spark.createDataFrame(data, ["features", "label"])

将DataFrame数据转换为TFRecord格式

output_dir = "hdfs:///tfrecord/output"

df.write.format("tfrecords").option("recordType", "Example").save(output_dir)