python-3.x - 以可变批量大小加载数据？

Question

我目前正在研究基于补丁的超分辨率。大多数论文将图像分成较小的补丁，然后将补丁用作模型的输入。我能够使用自定义数据加载器创建补丁。代码如下:

import torch.utils.data as data
from torchvision.transforms import CenterCrop, ToTensor, Compose, ToPILImage, Resize, RandomHorizontalFlip, RandomVerticalFlip
from os import listdir
from os.path import join
from PIL import Image
import random
import os
import numpy as np
import torch

def is_image_file(filename):
    return any(filename.endswith(extension) for extension in [".png", ".jpg", ".jpeg", ".bmp"])

class TrainDatasetFromFolder(data.Dataset):
    def __init__(self, dataset_dir, patch_size, is_gray, stride):
        super(TrainDatasetFromFolder, self).__init__()
        self.imageHrfilenames = []
        self.imageHrfilenames.extend(join(dataset_dir, x)
                                     for x in sorted(listdir(dataset_dir)) if is_image_file(x))
        self.is_gray = is_gray
        self.patchSize = patch_size
        self.stride = stride

    def _load_file(self, index):
        filename = self.imageHrfilenames[index]
        hr = Image.open(self.imageHrfilenames[index])
        downsizes = (1, 0.7, 0.45)
        downsize = 2
        w_ = int(hr.width * downsizes[downsize])
        h_ = int(hr.height * downsizes[downsize])
        aug = Compose([Resize([h_, w_], interpolation=Image.BICUBIC),
                       RandomHorizontalFlip(),
                       RandomVerticalFlip()])

        hr = aug(hr)
        rv = random.randint(0, 4)
        hr = hr.rotate(90*rv, expand=1)
        filename = os.path.splitext(os.path.split(filename)[-1])[0]
        return hr, filename

    def _patching(self, img):

        img = ToTensor()(img)
        LR_ = Compose([ToPILImage(), Resize(self.patchSize//2, interpolation=Image.BICUBIC), ToTensor()])

        HR_p, LR_p = [], []
        for i in range(0, img.shape[1] - self.patchSize, self.stride):
            for j in range(0, img.shape[2] - self.patchSize, self.stride):
                temp = img[:, i:i + self.patchSize, j:j + self.patchSize]
                HR_p += [temp]
                LR_p += [LR_(temp)]

        return torch.stack(LR_p),torch.stack(HR_p)

    def __getitem__(self, index):
        HR_, filename = self._load_file(index)
        LR_p, HR_p = self._patching(HR_)
        return LR_p, HR_p

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

假设批量大小为 1，它获取图像并给出大小为 [x,3,patchsize,patchsize] 的输出。当批量大小为 2 时，我将有两个不同大小的输出 [x,3,patchsize,patchsize] (例如图像 1 可能给出 [50,3,patchsize,patchsize]，图像 2 可能给出 [75,3,patchsize,patchsize] )。为了处理这个问题，需要一个自定义的整理函数，将这两个输出沿维度 0 堆叠。整理函数如下所示:

def my_collate(batch):
    data = torch.cat([item[0] for item in batch],dim = 0)
    target = torch.cat([item[1] for item in batch],dim = 0)

    return [data, target]

这个 collat​​e 函数沿 x 连接(从上面的示例中，我最终得到 [125,3,patchsize,pathsize]。出于训练目的，我需要使用一个 minibatch 大小的说 25. 是否有任何方法或任何函数可用于直接从数据加载器获取大小为 [25 , 3, patchsize, pathsize] 的输出，使用必要数量的图像作为输入数据加载器？

Answer 1

以下代码段适用于您的目的。

首先，我们定义一个ToyDataset，它接收一个张量列表(tensors)，0维变长。这类似于您的数据集返回的样本。

import torch
from torch.utils.data import Dataset
from torch.utils.data.sampler import RandomSampler

class ToyDataset(Dataset):
    def __init__(self, tensors):
        self.tensors = tensors

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

    def __len__(self):
        return len(tensors)

其次，我们定义了一个自定义数据加载器。创建数据集和数据加载器的常用 Pytorch 二分法大致如下:有一个带索引的 dataset，您可以将索引传递给它，它会从数据集中返回相关的样本。有一个产生索引的 sampler，有不同的策略来绘制索引，从而产生不同的采样器。 batch_sampler 使用采样器一次绘制多个索引(与 batch_size 指定的一样多)。有一个 dataloader 结合了采样器和数据集，让您可以迭代数据集，重要的是，数据加载器还拥有一个函数 (collat​​e_fn)，它指定如何从中检索多个样本应该组合使用来自 batch_sampler 的索引的数据集。对于您的用例，通常的 PyTorch 二分法效果不佳，因为我们不需要绘制固定数量的索引，而是需要绘制索引，直到与索引关联的对象超过我们想要的累积大小。这意味着我们需要立即检查对象并使用这些知识来决定是返回批处理还是保留绘图索引。这就是下面的自定义数据加载器所做的:

class CustomLoader(object):

    def __init__(self, dataset, my_bsz, drop_last=True):
        self.ds = dataset
        self.my_bsz = my_bsz
        self.drop_last = drop_last
        self.sampler = RandomSampler(dataset)

    def __iter__(self):
        batch = torch.Tensor()
        for idx in self.sampler:
            batch = torch.cat([batch, self.ds[idx]])
            while batch.size(0) >= self.my_bsz:
                if batch.size(0) == self.my_bsz:
                    yield batch
                    batch = torch.Tensor()
                else:
                    return_batch, batch = batch.split([self.my_bsz,batch.size(0)-self.my_bsz])
                    yield return_batch
        if batch.size(0) > 0 and not self.drop_last:
            yield batch

这里我们迭代数据集，在绘制索引并加载关联对象之后，我们将它连接到我们之前绘制的张量(batch)。我们一直这样做，直到达到所需的大小，这样我们就可以切出并生产一批。我们保留了 batch 中没有产生的行。因为可能会出现单个实例超出期望的batch_size的情况，所以我们使用while循环。

您可以修改这个最小的 CustomDataloader 以添加更多 PyTorch 数据加载器样式的功能。也不需要使用 RandomSampler 来绘制索引，其他的也同样可以。如果您的数据很大，例如通过使用列表并跟踪其张量的累积长度，也可以避免重复连接。

这是一个例子，证明它是有效的:

patch_size = 5
channels = 3
dim0sizes = torch.LongTensor(100).random_(1, 100)
data = torch.randn(size=(dim0sizes.sum(), channels, patch_size, patch_size))
tensors = torch.split(data, list(dim0sizes))

ds = ToyDataset(tensors)
dl = CustomLoader(ds, my_bsz=250, drop_last=False)
for i in dl:
    print(i.size(0))