使用 tf.data 读取 CSV 文件非常慢，改用 tfrecords ?

Question

我有很多CSV文件，每条记录包含~6000列。第一列是标签，其余列应被视为特征向量。我是Tensorflow的新手，我不知道如何将数据读取到具有所需格式的TensorflowDataset中。我当前正在运行以下代码：

DEFAULTS = []
n_features = 6170
for i in range(n_features+1):
DEFAULTS.append([0.0])
def parse_csv(line):
# line = line.replace('"', '')
columns = tf.decode_csv(line, record_defaults=DEFAULTS)  # take a line at a time
features = {'label': columns[-1], 'x': tf.stack(columns[:-1])}  # create a dictionary out of the features
labels = features.pop('label')  # define the label
return features, labels

def train_input_fn(data_file=sample_csv_file, batch_size=128):
"""Generate an input function for the Estimator."""
# Extract lines from input files using the Dataset API.
dataset = tf.data.TextLineDataset(data_file)
dataset = dataset.map(parse_csv)
dataset = dataset.shuffle(10000).repeat().batch(batch_size)
return dataset.make_one_shot_iterator().get_next()

每个 CSV 文件都有 ~10K 条记录。我试图在train_input_fn上做一个示例评估，作为labels = train_input_fn()[1].eval(session=sess).这需要 128 个标签，但大约需要2 分钟。

我是否使用了一些冗余操作，或者是否有更好的方法来执行此操作？

PS：我在Spark数据帧中有原始数据。因此，如果这可以使事情变得更快，我也可以使用 TFRecords。

Answer 1

你做对了。但更快的方法是使用TFRecords，如以下步骤所示：

1. 使用tf.python_io.TFRecordWriter： -- 要读取 csv 文件并将其写入 tfrecord 文件，如下所示：Tensorflow 从 csv 创建一个 tfrecords 文件。

2. 从 tfrecord 中读取： --

   def _parse_function(proto):
   f = {
   "features": tf.FixedLenSequenceFeature([], tf.float32, default_value=0.0, allow_missing=True),
   "label": tf.FixedLenSequenceFeature([], tf.float32, default_value=0.0, allow_missing=True)
   }
   parsed_features = tf.parse_single_example(proto, f)
   features = parsed_features["features"]
   label = parsed_features["label"]
   return features, label
   
   dataset = tf.data.TFRecordDataset(['csv.tfrecords'])
   dataset = dataset.map(_parse_function)
   dataset = dataset.shuffle(10000).repeat().batch(128)
   iterator = dataset.make_one_shot_iterator()
   features, label = iterator.get_next()
   

我在随机生成的 csv 上运行了这两个案例(csv vs tfrecords)。csv直读的10批次(每批128个样本)的总时间约为204s，而tfrecord的总时间约为0.22s。