PyTorch的Dataset和DataLoader

PyTorch的Dataset和DataLoader的例子

import torch
import numpy as np
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

# Dataset的子类必须重写__getitem__和__len__方法。
class TestDataset(Dataset):
    def __init__(self, src, tgt):
        self.src = src
        self.tgt = tgt
    def __getitem__(self, index):
        return self.src[index], self.tgt[index]
    def __len__(self):
        return len(self.src)

src_data = [
    np.array([1,2,3]),
    np.array([4,5,6]),
    np.array([7,8,9])
]

tgt_data = [
    np.array([11,12,13]),
    np.array([14,15,16]),
    np.array([17,18,19])
]

dataset = TestDataset(src_data, tgt_data)
dataload = DataLoader(dataset)

for src, tgt in dataload:
    print(src, tgt)

结果：

tensor([[1, 2, 3]]) tensor([[11, 12, 13]])
tensor([[4, 5, 6]]) tensor([[14, 15, 16]])
tensor([[7, 8, 9]]) tensor([[17, 18, 19]])

赏

谢谢你请我喝奶茶

支付宝

微信