Hello. I am recently trying to load tfrecords using pytorch. However, it seems that if I load tf.data.TFRecordDataset in pytorch datasets and use dataloader with num_workers > 0, the program won’t work properly. I am wondering if there is any better ways to load tfrecords or other better ways to store large scale datasets.
Here are the example codes:
class TestDataset(Dataset):
def __init__(self, record_path):
self.record_path = record_path
self.reader = tf.data.TFRecordDataset(self.record_path).map(decoder)
self._records_iter = self.reader.make_one_shot_iterator()
def __len__(self):
return 100
def _parser(self, img):
image_arr = np.frombuffer(img, dtype=np.uint8)
sample = torch.tensor(image_arr)
return sample
def __getitem__(self, item):
sample = next(self._records_iter).numpy()
return self._parser(sample)
dataset = TestDataset(0, path)
loader = DataLoader(dataset, batch_size=1, num_workers=1)
for i in loader:
print(i)