【深度好文】多任务模型中的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算法资讯.
