如何輕松地將 PyTorch 數據加載器轉換為 tf.Dataset？

Question

我們如何將pytorch數據加載器轉換為tf.Dataset ？

我發現了這個片段：-

def convert_pytorch_dataloader_to_tf_dataset(dataloader, batch_size, shuffle=True):
    dataset = tf.data.Dataset.from_generator(
        lambda: dataloader,
        output_types=(tf.float32, tf.float32),
        output_shapes=(tf.TensorShape([256, 512]), tf.TensorShape([2,]))
    )
    if shuffle:
        dataset = dataset.shuffle(buffer_size=len(dataloader.dataset))
    dataset = dataset.batch(batch_size)
    return dataset

但它根本不起作用。

是否有內置選項可以輕松地將數據tf.Dataset dataloaders 我有一個非常復雜的數據加載器，所以一個簡單的解決方案應該確保沒有錯誤:)

Answer 1

對於 h5py 格式的數據，您可以使用下面的腳本。 name_x 是 h5py 中的功能名稱，name_y 是標簽的文件名。 這種方法是 memory 有效的，您可以批量輸入數據。

class Generator(object):

def __init__(self,open_directory,batch_size,name_x,name_y):

    self.open_directory = open_directory

    data_f = h5py.File(open_directory, "r")

    self.x = data_f[name_x]
    self.y = data_f[name_y]

    if len(self.x.shape) == 4:
        self.shape_x = (None, self.x.shape[1], self.x.shape[2], self.x.shape[3])

    if len(self.x.shape) == 3:
        self.shape_x = (None, self.x.shape[1], self.x.shape[2])

    if len(self.y.shape) == 4:
        self.shape_y = (None, self.y.shape[1], self.y.shape[2], self.y.shape[3])

    if len(self.y.shape) == 3:
        self.shape_y = (None, self.y.shape[1], self.y.shape[2])

    self.num_samples = self.x.shape[0]
    self.batch_size = batch_size
    self.epoch_size = self.num_samples//self.batch_size+1*(self.num_samples % self.batch_size != 0)

    self.pointer = 0
    self.sample_nums = np.arange(0, self.num_samples)
    np.random.shuffle(self.sample_nums)


def data_generator(self):

    for batch_num in range(self.epoch_size):

        x = []
        y = []

        for elem_num in range(self.batch_size):

            sample_num = self.sample_nums[self.pointer]

            x += [self.x[sample_num]]
            y += [self.y[sample_num]]

            self.pointer += 1

            if self.pointer == self.num_samples:
                self.pointer = 0
                np.random.shuffle(self.sample_nums)
                break

        x = np.array(x,
                     dtype=np.float32)
        y = np.array(y,
                     dtype=np.float32)

        yield x, y

def get_dataset(self):
    dataset = tf.data.Dataset.from_generator(self.data_generator,
                                             output_types=(tf.float32,
                                                           tf.float32),
                                             output_shapes=(tf.TensorShape(self.shape_x),
                                                            tf.TensorShape(self.shape_y)))
    dataset = dataset.prefetch(1)

    return dataset