如何创建 torch.utils.data.Dataset 并将其导入到 torch.utils.data.DataLoader答案

【问题标题】：How to create a torch.utils.data.Dataset and import it into a torch.utils.data.DataLoader如何创建 torch.utils.data.Dataset 并将其导入到 torch.utils.data.DataLoader
【发布时间】：2021-11-26 16:35:50
【问题描述】：

我想从带有torch.utils.data.Dataset 的CSV 文件中导入数据，这样我就可以使用torch.utils.data.DataLoader 来处理它。数据不属于torchvision，来自我的电脑。谷歌上似乎没有解决方案。如果您能给我一些建议，我将非常感谢。

【问题讨论】：

标签： python torchvision

【解决方案1】：

如果您已经拥有 csv 文件，则可以使用 pandas 轻松完成此操作。

import pandas as pd
my_dataframe = pd.read_csv("path/to/file.csv")

现在，您可以访问 csv 文件中的数据。如果你想使用 pytorch torch.utils.data.DataLoader，你还需要一个 torch.utils.data.Dataset。

根据您使用的数据类型，数据集可能看起来非常不同。如果您正在处理 csv 中的图像路径和标签，请查看我曾经用于 torchvision.models.resnet50() 的这个数据集：

from torch.utils.data import Dataset
from PIL import Image
from torchvision import models, transforms
import cv2

class createDataset(Dataset):
    def __init__(self, dataframe):
        self.dataframe = dataframe
        self.transform = transforms.Compose([transforms.ToTensor()])

    def __len__(self):
        return self.dataframe.shape[0]
        
    def __getitem__(self, index):
        image = self.dataframe.iloc[index]["Name_of_imagepath_column"]
        image = cv2.imread(image)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = Image.fromarray(image)    
        image = self.transform(image)
        label = self.dataframe.iloc[index]["Name_of_label_column"]
        return {"image": image , "targets": torch.tensor(label, dtype=torch.long)}

标签/目标是可选的，仅在我的项目中是必需的。

现在您可以将您的 pandas 数据框传递给 Dataset 类，如下所示：

my_dataset = createDataset(dataframe = my_dataframe)

现在可以将此数据集传递给torch.utils.data.DataLoader 并创建您的数据加载器：

from torch.utils.data import DataLoader

my_dataloader= DataLoader(dataset=my_dataset)

有关 Dataloader 的更多选项，例如批量大小和随机播放，请查看 Pytorch DataLoader docs

【讨论】：