溫馨提示×

pytorch中怎么創(chuàng)建自己的數(shù)據(jù)集

小億
95
2024-04-08 13:46:40

在PyTorch中,可以通過繼承torch.utils.data.Dataset類來創(chuàng)建自己的數(shù)據(jù)集。以下是一個簡單的示例代碼:

import torch
from torch.utils.data import Dataset

class CustomDataset(Dataset):
    def __init__(self, data):
        self.data = data
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        sample = self.data[idx]
        return sample

# 創(chuàng)建自己的數(shù)據(jù)集
data = [1, 2, 3, 4, 5]
custom_dataset = CustomDataset(data)

# 創(chuàng)建數(shù)據(jù)加載器
data_loader = torch.utils.data.DataLoader(custom_dataset, batch_size=2, shuffle=True)

# 遍歷數(shù)據(jù)加載器
for batch in data_loader:
    print(batch)

在上面的示例中,首先定義了一個自定義的數(shù)據(jù)集CustomDataset,該數(shù)據(jù)集繼承自torch.utils.data.Dataset類,并實現(xiàn)了__init__、__len____getitem__方法。然后創(chuàng)建了一個包含一些數(shù)據(jù)的實例data,并使用它來實例化CustomDataset類得到custom_dataset。最后,使用DataLoader將自定義數(shù)據(jù)集包裝成數(shù)據(jù)加載器,并遍歷數(shù)據(jù)加載器來獲取數(shù)據(jù)。

0