以可变批量大小加载数据？答案

【问题标题】：Data loading with variable batch size?以可变批量大小加载数据？
【发布时间】：2019-01-06 04:37:08
【问题描述】：

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

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]

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

【问题讨论】：

所以你希望你的 data_loader 总是返回一个大小为 [25,3, patch size, patchsize] 的输出，不管需要多少图像（即你希望它加载尽可能多的图像需要生成上述大小的输出）？
请。有什么办法吗？
一种简单的方法可能是定义一个自定义采样器，该采样器在达到 25 阈值之前不会加载图像，然后将其返回，同时保留其余的用于下一次迭代。我大概可以给你写一个简单的蓝图。
我想过使用采样器，但我不知道它的结构。如果您能提供蓝图，我将非常高兴。我可以以此为基础。

标签： python-3.x image-processing pytorch

【解决方案1】：

以下代码 sn-p 适用于您的目的。

首先，我们定义一个ToyDataset，它接收variable length in dimension 0的张量列表（tensors）。这类似于您的数据集返回的样本。

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))

【讨论】：

【解决方案2】：

（相关，但不完全在主题中）

对于批量大小调整，您可以使用this repo 中示例的代码。它是为了不同的目的（最大化 GPU 内存使用）而实现的，但它并不难转化为您的问题。

代码进行批量适配和批量欺骗。

【讨论】：

【解决方案3】：

为了改进之前的答案，我找到了一个repo，它使用DataManger 来实现不同的patch size 和batch size。它基本上是用不同的设置启动不同的数据加载器，set_epoch 函数用于为给定的时期设置适当的数据加载器。

【讨论】：