将tf.dataset转换为PyTorch数据集？

Question

在这个项目中，所有数据都经过预处理，并准备好作为tensorflow数据集，如下所示：

<MapDataset shapes: {input_ids: (128,), input_mask: (128,), label_ids: (), segment_ids: (128,)}, types: {input_ids: tf.int64, input_mask: tf.int64, label_ids: tf.int64, segment_ids: tf.int64}>

我拥有的脚本在PyTorch中，并接受一个Dataset对象，如下所示：

Dataset({
    features: ['attention_mask', 'input_ids', 'label', 'sentence', 'token_type_ids'],
    num_rows: 12
})

有什么办法把其中一个转换成另一个吗？我对这两个API都很陌生，所以我不太确定它们是如何工作的？我是否可以使用dict将其中一个转换成另一个呢？

谢谢

Answer 1

我使用tfds.as_numpy(dataset)作为我的模型培训的数据处理程序。为了转换传递给我的模型的数据，我在模型的前向函数中使用了torch.as_tensor(data, device=<device>)。

import tensorflow_datasets as tfds
import torch.nn as nn

def train_dataloader(batch_size):
    return tfds.as_numpy(tfds.load('mnist').batch(batch_size))

class Model(nn.Module):
    def forward(self, x):
        x = torch.as_tensor(x, device='cuda')
        ...