Python Script: Converts .tsv to TFRecord

Question

The python code below converts a tab spaced values file (.tsv) into a TFRecord in hopes that this approach expedites future time spent loading the data. Is the performance of this task correct and optimal, specifically in terms of run-time?

The .tsv file is of the following format, including spaces to make it more readable (SEMEVAL STS'17 training dataset): float_value_from_0_to_5 \t string1 \t string2. The desired TFRecord's entries should optimally store the previous values as float rate, string sent1, and string sent2 for proper implementation with any arbitrary Neural Network that will embed the words once loaded from the TFRecord.

import tensorflow as tf
import nltk
import os
import csv

def tsv_to_tfrec(file_path):
  with open(file_path, 'rb') as tsv, tf.python_io.TFRecordWriter(
        os.path.splitext(file_path)[0] + ".tfrecords") as tfrec:
    read_tsv = csv.reader(tsv, delimiter = '\t', quoting = csv.QUOTE_NONE)

    for row in read_tsv:
      if len(row) != 3:
        print("\nERROR in # of indicies: ", len(row))
        for i in range(len(row)):
          print(i, " =", row[i])
      else:
        example = tf.train.Example(features = tf.train.Features(feature = {
            "rate": tf.train.Feature(float_list=tf.train.FloatList(
                value = [float(row[0])])),
            "sent1": tf.train.Feature(bytes_list=tf.train.BytesList(
                value = nltk.word_tokenize(row[1]))),
            "sent2": tf.train.Feature(bytes_list=tf.train.BytesList(
                value = nltk.word_tokenize(row[2])))
            }))

        tfrec.write(example.SerializeToString())

Answer 1

I am no expert in TensorFlow (even though started to grasp), but here are some things I noticed:

• if this is a proper valid TSV file, you should not be worried about checking the length of each row - each row should have the same amount of columns

• you may also try to use TextLineReader to read features from a csv file like noted here, this may lead to "unpacking" the features and avoiding getting them by index:

  rate, sent1, sent2 = tf.decode_csv(value)
  

• if, for some reason, TextLineReader is not applicable, you may unpack in the for loop:

  for rate, sent1, sent2 in read_tsv:
  

• you can also look into applying some micro-optimizations, like avoiding name lookups in the loop by pre-lookuping things before the loop - e.g. setting Example = tf.train.Example and then using Example in place of tf.train.Example
• there is also a possibility of that classic trading off memory for the speed. Depending on how unique your features are, you, for instance, can "cache" the word tokenization results in a dictionary to avoid re-tokenizing the same sentences over again

Also check out the TensorFlow Performance Guide and Performance Tips.

There are also some minor PEP8 code style violations, like extra spaces around = in the keyword argument definitions - run your code through a static code analyzer like flake8 or pylint.

Answer 2

alecxe's answer's intent was to provide suggestions, I would like to provide the actual revised code to depict his suggestions in practice. The below code includes the implementation of all of his suggestions, but using unpacking of variables in the for loop rather than a TextLineReader. This is because a TextLineReader returns a tensor, and I do not know of an easy way to go from tensor to TFRecord, especially since the sentences need to be tokenized into a list of words. Also, this code assumes that the .tsv file is of proper format, and expects a very specific format (float, string, string). If errors in .tsv format are possible, you will have to modify the code for error handling.

Another thing to note, when it comes to the 'micro-optimizations' for doing name look-ups once, it may be better or equivalent performance to import the specific name look-up in the function and set its namespace as 'name'. For example with Example: import tf.train.Example as Example. The reason to keep it in the function, is to avoid muddying the namespace in the whole file. If you do not care about that, then just put it with the rest of the imports at the top. I am uncertain, if this is a better approach so I went with what was more readable and understandable.

The use of sentence_dict is to implement the trade off between memory for speed. If you're feeling especially ambitious, which I was not, you could create a personal function for sentence/string tokenization based on nltk's that will also store a word_dict so that new sentences with previously seen words will have constant time look up, but this may cause more time to compute due to each word needing to be checked before tokenization. Nonetheless, it is an idea that may be worth exploring.

I was doubtful on the speed boost from the sentence_dict due to not knowing if checking every sentence would take less run-time than just tokenizing all the words, but after comparing the code with and without the sentence_dict, I found a noteable decrease in run-time (by about a tenth of a second). This was using the SEMEVAL STS 2017 training data for the task of English to English semantic matching.

import os
import csv
import tensorflow as tf
import nltk

def tsv_to_tfrec(filename):
    """
    Creates tfrecord equivalent to .tsv provided by filename. Assumes proper
    .tsv with equal columns in each row.

    @param filename: filename of .tsv to be converted to TFRecord
    """

    with open(filename, 'rb') as tsv, tf.python_io.TFRecordWriter(
        os.path.splitext(filename)[0] + ".tfrecords") as tfrec:

        read_tsv = csv.reader(tsv, delimiter='\t', quoting=csv.QUOTE_NONE)

        # Trading off memory for speed
        sentence_dict = {}

        # Micro-Optimize by name look-up once
        Example = tf.train.Example
        Feature = tf.train.Feature
        Features = tf.train.Features
        FloatList = tf.train.FloatList
        BytesList = tf.train.BytesList

        for rate, sent1, sent2 in read_tsv:
            if sent1 not in sentence_dict:
                sentence_dict[sent1] = nltk.word_tokenize(sent1)

            if sent2 not in sentence_dict:
                sentence_dict[sent2] = nltk.word_tokenize(sent2)

            example = Example(features=Features(feature={
                "rate": Feature(float_list=FloatList(
                    value=[float(rate)])),
                "sent1": Feature(bytes_list=BytesList(
                    value=sentence_dict[sent1])),
                "sent2": Feature(bytes_list=BytesList(
                    value=sentence_dict[sent2]))
                }))

            tfrec.write(example.SerializeToString())