从pandas加载图像数据到pytorch

Question

我完全不懂pytorch，之前是在keras和fastai上工作的。 现在正在尝试一个图像回归任务，并且挑战是我必须从pandas数据帧中加载数据。 数据帧结构：

ID   Path   Score
fig1  /folder/fig1.jpg  2
fig2  /folder/fig2.jpg  3
.....

我之前曾经从文件夹直接加载图像到PyTorch中，因为那是一个简单的分类任务，但现在卡住了。

我查阅了PyTorch论坛的相关资料，但并没有完全理解如何实施该功能。非常感谢您的帮助！

- Akshat Shreemali

Answer 1

数据集

您需要使用torch.utils.data.Dataset结构来定义它。以下是在普通的pytorch中如何执行此操作（我使用pillow加载图像，并使用torchvision将它们转换为torch.Tensor对象）:

import torch
import torchvision
from PIL import Image


class MyDataset(torch.utils.data.Dataset):
    def __init__(self, dataframe):
        self.dataframe = dataframe

    def __len__(self):
        return len(self.dataframe)

    def __getitem__(self, index):
        row = self.dataframe.iloc[index]
        return (
            torchvision.transforms.functional.to_tensor(Image.open(row["Path"])),
            row["Score"],
        )

dataset = MyDataset(dataframe)

或者，您可以使用torchdata（免责声明：我是作者，这是自荐...），它允许您像这样解耦Path和Scores：

import torchvision
from PIL import Image

import torchdata


class ImageDataset(torchdata.datasets.FilesDataset):
    def __getitem__(self, index):
        return Image.open(self.files[index])


class Labels(torchdata.Dataset):
    def __init__(self, scores):
        super().__init__()
        self.scores = scores

    def __len__(self):
        return len(self.scores)

    def __getitem__(self, index):
        return self.scores[index]

# to_numpy for convenience
# I assume all your images are in /folder and have *.jpg extension
dataset = ImageDataset.from_folder("/folder", regex="*.jpg").map(
    torchvision.transforms.ToTensor()
) | Labels(dataframe["Score"].to_numpy())

（或者您可以像在常规的pytorch中那样实现，但是要从torchdata.Dataset继承并在构造函数中调用super().__init__()）。

torchdata允许您通过.map轻松缓存图像或应用其他转换，如下所示，请查看Github仓库以获取更多信息，或在评论中提问。

DataLoader

无论哪种方式，您都应该将数据集包装在torch.utils.data.DataLoader中，以创建批次并进行迭代处理，如下所示：

dataloader = torch.utils.data.DataLoader(dataset, batch_size=64, shuffle=True)

for images, scores in dataloader:
    # Rest of your code to train neural network or smth
    ...

在循环中，可随意处理这些图像和分数。

- Szymon Maszke