Package: tfrecord
TFRecord reader and writerThis library allows reading and writing tfrecord files efficiently in python. The library also provides an IterableDataset reader of tfrecord files for PyTorch. Currently uncompressed and compressed gzip TFRecords are supported.Installationpip3 install tfrecordUsageIt's recommended to create an index file for each TFRecord file. Index file must be provided when using multiple workers, otherwise the loader may return duplicate records. You can create an index file for an individual tfrecord file with this utility program:python3 -m tfrecord.tools.tfrecord2idx <tfrecord path> <index path>To create ".tfidnex" files for all ".tfrecord" files in a directory run:tfrecord2idx <data dir>Reading & Writing tf.train.ExampleReading tf.Example records in PyTorchUse TFRecordDataset to read TFRecord files in PyTorch.importtorchfromtfrecord.torch.datasetimportTFRecordDatasettfrecord_path="/tmp/data.tfrecord"index_path=Nonedescription={"image":"byte","label":"float"}dataset=TFRecordDataset(tfrecord_path,index_path,description)loader=torch.utils.data.DataLoader(dataset,batch_size=32)data=next(iter(loader))print(data)Use MultiTFRecordDataset to read multiple TFRecord files. This class samples from given tfrecord files with given probability.importtorchfromtfrecord.torch.datasetimportMultiTFRecordDatasettfrecord_pattern="/tmp/{}.tfrecord"index_pattern="/tmp/{}.index"splits={"dataset1":0.8,"dataset2":0.2,}description={"image":"byte","label":"int"}dataset=MultiTFRecordDataset(tfrecord_pattern,index_pattern,splits,description)loader=torch.utils.data.DataLoader(dataset,batch_size=32)data=next(iter(loader))print(data)Infinite and finite PyTorch datasetBy default,MultiTFRecordDatasetis infinite, meaning that it samples the data forever. You can make it finite by providing the appropriate flagdataset = MultiTFRecordDataset(..., infinite=False)Shuffling the dataBoth TFRecordDataset and MultiTFRecordDataset automatically shuffle the data when you provide a queue size.dataset = TFRecordDataset(..., shuffle_queue_size=1024)Transforming input dataYou can optionally pass a function astransformargument to perform post processing of features before returning.
This can for example be used to decode images or normalize colors to a certain range or pad variable length sequence.importtfrecordimportcv2defdecode_image(features):# get BGR image from bytesfeatures["image"]=cv2.imdecode(features["image"],-1)returnfeaturesdescription={"image":"bytes",}dataset=tfrecord.torch.TFRecordDataset("/tmp/data.tfrecord",index_path=None,description=description,transform=decode_image)data=next(iter(dataset))print(data)Writing tf.Example records in Pythonimporttfrecordwriter=tfrecord.TFRecordWriter("/tmp/data.tfrecord")writer.write({"image":(image_bytes,"byte"),"label":(label,"float"),"index":(index,"int")})writer.close()Reading tf.Example records in Pythonimporttfrecordloader=tfrecord.tfrecord_loader("/tmp/data.tfrecord",None,{"image":"byte","label":"float","index":"int"})forrecordinloader:print(record["label"])Reading & Writing tf.train.SequenceExampleSequenceExamples can be read and written using the same methods shown above with an extra argument
(sequence_descriptionfor reading andsequence_datumfor writing) which cause the respective
read/write functions to treat the data as a SequenceExample.Writing SequenceExamples to fileimporttfrecordwriter=tfrecord.TFRecordWriter("/tmp/data.tfrecord")writer.write({'length':(3,'int'),'label':(1,'int')},{'tokens':([[0,0,1],[0,1,0],[1,0,0]],'int'),'seq_labels':([0,1,1],'int')})writer.write({'length':(3,'int'),'label':(1,'int')},{'tokens':([[0,0,1],[1,0,0]],'int'),'seq_labels':([0,1],'int')})writer.close()Reading SequenceExamples in pythonReading from a SequenceExample yeilds a tuple containing two elements.importtfrecordcontext_description={"length":"int","label":"int"}sequence_description={"tokens":"int","seq_labels":"int"}loader=tfrecord.tfrecord_loader("/tmp/data.tfrecord",None,context_description,sequence_description=sequence_description)forcontext,sequence_featsinloader:print(context["label"])print(sequence_feats["seq_labels"])Read SequenceExamples in PyTorchAs described in the section onTransforming Input, one can pass a function as thetransformargument to
perform post processing of features. This should be used especially for the sequence features as these are
variable length sequence and need to be padded out before being batched.importtorchimportnumpyasnpfromtfrecord.torch.datasetimportTFRecordDatasetPAD_WIDTH=5defpad_sequence_feats(data):context,features=datafork,vinfeatures.items():features[k]=np.pad(v,((0,PAD_WIDTH-len(v)),(0,0)),'constant')return(context,features)context_description={"length":"int","label":"int"}sequence_description={"tokens":"int ","seq_labels":"int"}dataset=TFRecordDataset("/tmp/data.tfrecord",index_path=None,description=context_description,transform=pad_sequence_feats,sequence_description=sequence_description)loader=torch.utils.data.DataLoader(dataset,batch_size=32)data=next(iter(loader))print(data)Alternatively, you could choose to implement a customcollate_fnin order to assemble the batch,
for example, to perform dynamic padding.importtorchimportnumpyasnpfromtfrecord.torch.datasetimportTFRecordDatasetdefcollate_fn(batch):fromtorch.utils.data._utilsimportcollatefromtorch.nn.utilsimportrnncontext,feats=zip(*batch)feats_={k:[torch.Tensor(d[k])fordinfeats]forkinfeats[0]}return(collate.default_collate(context),{k:rnn.pad_sequence(f,True)for(k,f)infeats_.items()})context_description={"length":"int","label":"int"}sequence_description={"tokens":"int ","seq_labels":"int"}dataset=TFRecordDataset("/tmp/data.tfrecord",index_path=None,description=context_description,transform=pad_sequence_feats,sequence_description=sequence_description)loader=torch.utils.data.DataLoader(dataset,batch_size=32,collate_fn=collate_fn)data=next(iter(loader))print(data)