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

Question

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

<MapDataset形状：{input_ids：（128，），input_mask：（128，），label_ids：（），segment_ids：（128，）}，类型：{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 都很陌生，所以不太清楚它们是如何工作的。我可以将其中一种转换为另一种吗？

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')
        ...