使用可变批量大小加载数据？

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

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，它获取图像并给出 size 的输出[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]

这个整理函数沿着 x 连接（从上面的例子中，我终于得到[125,3,patchsize,pathsize]。出于训练目的，我需要使用 25 的小批量大小来训练模型。是否有任何方法或函数可以用来直接获得大小的输出[25 , 3, patchsize, pathsize]直接从数据加载器使用必要数量的图像作为数据加载器的输入？

以下代码片段可满足您的目的。

首先，我们定义一个 ToyDataset，它接受张量列表（tensors) of variable length in dimension 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它结合了采样器和数据集，让您可以迭代数据集，重要的是数据加载器还拥有一个函数（collate_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 loop.

您可以修改这个最小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))