I made a working version of this with IterableDataset but it was unusably slow. It was not the TfRecord reader’s fault but IterableDataset, it seemed like it didn’t use queuing at all.
To provide more context on why I want to use something like TFRecords, I work with audio files that are all very small, so a large dataset ends up being 90MM files totaling 5Tb. That becomes very unwieldy to move around, so serializing the encoded FLACs into a protobuf format and chunking them makes them much easier to move around and read from a network.
Minimal example:
import tensorflow as tf
import torch
import random

class TfRecordDataset(torch.utils.data.IterableDataset):

def __init__(self, tfrecord_path):
        "tfrecord_path contains tf record files"""
        self.buffer_size = buffer_size
        self.tfrecord_list = sorted(glob.glob(tfrecord_path))
        super(TfRecordDataset, self).__init__()

def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        if worker_info is None:
            tfrecord_list = self.tfrecord_list
        else:
            worker_id = worker_info.id
            num_workers = worker_info.num_workers
            tfrecord_list = self.tfrecord_list[worker_id::num_workers]
            random.seed(worker_info.seed)

        def tf_record_iterator():
            for filename in tfrecord_list:
                    record_iterator = tf.python_io.tf_record_iterator(filename,
                                            tf.python_io.TFRecordCompressionType.NONE)
                    for string_record in record_iterator:
                        example = tf.train.Example()
                        example.ParseFromString(string_record)
                        do_stuff_with_example(example)
                        yield example

        return tf_record_iterator()

Does something look obviously wrong this implementation to the PyTorch people?