pyspark将DataFrame写入tfrecord 和tf中tfrecord的解析

做特征工程时，常需要利用pyspark 将数据处理后保存为tfrecord以供后续的模型训练。本文给一个保存为tfrecord 和对tfrecord进行解析的案例。

一 、spark将DataFrame写入tfrecord 

用spark读取数据经过处理后，保存为pyspark的dataframe形式。

def tfrecord(data,dt_begin,dt_end,is_train=True):

        from pyspark.sql import functions as funcs

        if is_train: # 如果是train，将数据打乱下（在写tfrecord时就进行) 

                data = data.orderBy(funcs.rand())

                tag = 'train_{0}_{1}_{2}'.format(VERSION, dt_begin, dt_end)

        else:

                tag = 'test_{0}_{1}_{2}'.format(VERSION, dt_begin, dt_end)

        sample_count = data.cache().count()

        # fill  对各数值列进行缺失值的进一步填充

        data = data.fillna(0.0, numerical_features) # 也可以用字典的形式对各指定列以指定值填充

        # select ，从dataframe中选取要保存为tfrecord的数据。主要包括labels列(多目标可能有不止一个label) 和 types_features特征列。 

        cols = [funcs.array(*[funcs.col(col).cast('double') for col in labels]).alias('labels')]  # labels 为double类型的

        cols += [funcs.array(*(types_features)).alias('features')]  # types_features 特征列此时也转换为了数值类型了

        # write ,将需要保存的列写入为tfrecord 

        data = data.select(cols).persist()

        tag = '{0}_{1}_{2}'.format(version, beg, end) ## 保存的文件名(带上版本号容易区分) 

        data.coalesce(3).write.mode('overwrite').format('tfrecords').option('recordType', 'Example').save(path + tag)

        data.unpersist()

二、对应的tfrecord解析

对于保存的数据进行tfrecord解析，需要注意格式要对应。 

def input_fn(path, params, is_train=True):
    def f_parse(record):
        data = tf.parse_example(record, {
            'labels': tf.FixedLenFeature([labels列中label长度值即对应上面len(labels)的值], dtype=tf.float32),
            'features': tf.FixedLenFeature([特征列中的各个特征值个数的长度(多值被碾平)即对应上面len(types_features)的值], dtype=tf.float32)
        })
        return data

    with tf.variable_scope('tf_reader'):
        filenames = hdfs.connect().ls(tfrord_存储的hdfs_path)
        fileset = tf.data.Dataset.from_tensor_slices(filenames)
        if is_train:  # shuffle file name
            fileset = fileset.shuffle(len(filenames)).repeat()
        dataset = tf.data.TFRecordDataset(fileset, buffer_size=100000, num_parallel_reads=4)
        if is_train:  # shuffle data
            dataset = dataset.shuffle(10000)
        return dataset.batch(params.batch_size).map(f_parse, 4).prefetch(10000)

返回结果返回一个batch的数据，后续供模型训练。