【深度好文】多任务模型中的DataLoader实现

对于多任务学习multi-task-learning(MTL)问题，经常会要求特定的训练过程，比如数据处理，模型结构和性能评估函数．本文主要针对数据处理部分进行展开，主要针对多个标注好的数据集如何来训练一个多任务模型．

本文主要从两个方面进行展开：
１．将两个或多个dataset组合成pytorch中的一个Dataset．这个dataset将会作为pytorch中Dataloader的输入．
２．修改batch产生过程，以确保在第一个batch中产生第一个任务的数据，在第二个batch中产生下一个任务的数据．

为了简单处理，我们将以两个dataset作为例子来讲述．通常来说，dataset的数目以及data的类型不会对我们的整体方案带来太大影响．一个pytorch的Dataset需要实现　＿＿getitem__()函数．这个函数的作用为预取数据并且为给定index准备数据．

第一节　定义dataset

首先，我们先来定义两个dummy dataset,如下所示：

import torch
from torch.utils.data.dataset import ConcatDataset


class MyFirstDataset(torch.utils.data.Dataset):
    def __init__(self):
        # dummy dataset
        self.samples = torch.cat((-torch.ones(5), torch.ones(5)))

    def __getitem__(self, index):
        # change this to your samples fetching logic
        return self.samples[index]

    def __len__(self):
        # change this to return number of samples in your dataset
        return self.samples.shape[0]


class MySecondDataset(torch.utils.data.Dataset):
    def __init__(self):
        # dummy dataset
        self.samples = torch.cat((torch.ones(50) * 5, torch.ones(5) * -5))

    def __getitem__(self, index):
        # change this to your samples fetching logic
        return self.samples[index]

    def __len__(self):
        # change this to return number of samples in your dataset
        return self.samples.shape[0]


first_dataset = MyFirstDataset()
second_dataset = MySecondDataset()
concat_dataset = ConcatDataset([first_dataset, second_dataset])

上述代码中，我们定义了两个dataset,其中第一个dataset长度为10,其中前５个sample为-1，后５个sample为１；其中第二个dataset长度为55，其中前50个sample为５，后５个sample为-5．上述数据集仅仅为了说明方便．在实际应用中，我们应该会同时拥有sample和label，当然我们也可能会从一个目录或者数据库中读取数据，但是上面简单的dataset足够帮助我们来了解整个实现流程．

第二节 定义dataloader

接着我们来定义Dataloader，这里我们使用pytorch中的concat_data来实现两个dataset的合并．
代码如下：

batch_size = 8

# basic dataloader
dataloader = torch.utils.data.DataLoader(dataset=concat_dataset,
                                         batch_size=batch_size,
                                         shuffle=True,
                                         drop_last=True
                                         )

for inputs in dataloader:
    print(inputs)

运行结果如下：

tensor([ 5.,  5.,  5.,  5., -1.,  5.,  5.,  5.])
tensor([ 5.,  1., -1., -1.,  5.,  5.,  5., -5.])
tensor([5., 5., 5., 5., 5., 5., 5., 5.])
tensor([ 5., -5., -5.,  5.,  5.,  5.,  5.,  5.])
tensor([-1.,  5., -1.,  5.,  5.,  5.,  5.,  5.])
tensor([ 5.,  5., -5.,  5.,  5.,  5.,  5.,  1.])
tensor([5., 5., 5., 5., 1., 5., 5., 5.])
tensor([ 5.,  1.,  5., -5.,  5.,  5.,  1.,  5.])

对于我们的concat_dataset来说，每个batch有８个sample．每个sample的次序是随机的．

第三节 定义sampler

到现在为止，上述实现都很简单直接．上述dataset被合并成一个dataset，并且sample都是从原先dataset中随机挑选组合成batch的．现在让我们来写控制每个batch中的sample来源．我们预期达到的目的在每一个batch中，数据仅来自一个task的dataset,在下一个batch中进行切换．此时我们需要自己定义sample,其代码实现如下：

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


class BatchSchedulerSampler(torch.utils.data.sampler.Sampler):
    """
    iterate over tasks and provide a random batch per task in each mini-batch
    """
    def __init__(self, dataset, batch_size):
        self.dataset = dataset
        self.batch_size = batch_size
        self.number_of_datasets = len(dataset.datasets)
        self.largest_dataset_size = max([len(cur_dataset.samples) for cur_dataset in dataset.datasets])

    def __len__(self):
        return self.batch_size * math.ceil(self.largest_dataset_size / self.batch_size) * len(self.dataset.datasets)

    def __iter__(self):
        samplers_list = []
        sampler_iterators = []
        for dataset_idx in range(self.number_of_datasets):
            cur_dataset = self.dataset.datasets[dataset_idx]
            sampler = RandomSampler(cur_dataset)
            samplers_list.append(sampler)
            cur_sampler_iterator = sampler.__iter__()
            sampler_iterators.append(cur_sampler_iterator)

        push_index_val = [0] + self.dataset.cumulative_sizes[:-1]
        step = self.batch_size * self.number_of_datasets
        samples_to_grab = self.batch_size
        # for this case we want to get all samples in dataset, this force us to resample from the smaller datasets
        epoch_samples = self.largest_dataset_size * self.number_of_datasets

        final_samples_list = []  # this is a list of indexes from the combined dataset
        for _ in range(0, epoch_samples, step):
            for i in range(self.number_of_datasets):
                cur_batch_sampler = sampler_iterators[i]
                cur_samples = []
                for _ in range(samples_to_grab):
                    try:
                        cur_sample_org = cur_batch_sampler.__next__()
                        cur_sample = cur_sample_org + push_index_val[i]
                        cur_samples.append(cur_sample)
                    except StopIteration:
                        # got to the end of iterator - restart the iterator and continue to get samples
                        # until reaching "epoch_samples"
                        sampler_iterators[i] = samplers_list[i].__iter__()
                        cur_batch_sampler = sampler_iterators[i]
                        cur_sample_org = cur_batch_sampler.__next__()
                        cur_sample = cur_sample_org + push_index_val[i]
                        cur_samples.append(cur_sample)
                final_samples_list.extend(cur_samples)

        return iter(final_samples_list)

上述定义了一个BatchSchedulerSampler类，实现了一个新的sampler iterator.首先，通过为每一个单独的dataset创建RandomSampler；接着，在每一个dataset iter中获取对应的sample index；最后，创建新的sample index list．这里我们使用batchsize=８,那么我们将会从每个dataset中预取８个samples.
接着我们来测试上述sampler,代码如下：

import torch
from multi_task_batch_scheduler import BatchSchedulerSampler

batch_size = 8

# dataloader with BatchSchedulerSampler
dataloader = torch.utils.data.DataLoader(dataset=concat_dataset,
                                         sampler=BatchSchedulerSampler(dataset=concat_dataset,
                                                                       batch_size=batch_size),
                                         batch_size=batch_size,
                                         shuffle=False)

for inputs in dataloader:
    print(inputs)

运行结果如下：

tensor([ 1., -1.,  1.,  1., -1., -1., -1.,  1.])
tensor([ 5.,  5.,  5.,  5.,  5., -5.,  5., -5.])
tensor([ 1., -1., -1., -1., -1.,  1.,  1.,  1.])
tensor([5., 5., 5., 5., 5., 5., 5., 5.])
tensor([ 1.,  1., -1., -1.,  1.,  1.,  1.,  1.])
tensor([5., 5., 5., 5., 5., 5., 5., 5.])
tensor([-1.,  1., -1., -1., -1., -1.,  1., -1.])
tensor([-5.,  5.,  5.,  5.,  5.,  5.,  5.,  5.])
tensor([ 1., -1.,  1., -1., -1.,  1., -1.,  1.])
tensor([ 5., -5.,  5.,  5.,  5.,  5.,  5.,  5.])
tensor([-1., -1.,  1., -1.,  1., -1., -1.,  1.])
tensor([ 5.,  5.,  5., -5.,  5.,  5.,  5.,  5.])
tensor([ 1.,  1., -1., -1.,  1.,  1.,  1.,  1.])
tensor([5., 5., 5., 5., 5., 5., 5., 5.])

Wow,综上，我们实现了每一个minibatch仅从一个dataset中取数据的功能，并且下一个minibatch从不同任务的dataset中取batch.

参考：链接

关注公众号《AI算法之道》，获取更多AI算法资讯.