如果训练数据很多的话,就不能够使用pytorch提供的Dataset类来实现数据的读取了,因为这个类会在训练前生成一个跟数据量成比例的缓存,如果数据量达到100亿个数据点以上的话,占用的内存可能达到1TB以上,而且代码的运行效率也会显著下降。
查了一下google,发现有很多人有类似的需求,但是没有人提出了解决方案。在pytorch的issue栏目中可以看到类似的需求,最后也是没有提出解决方案。
当数据量太大时,我们需要使用IterableDataset来从硬盘直接读取数据。可以参考如下的方式进行实现( 效率上应该可以与tfrecord等数据存储结构 持平,不需要过度优化即可打满GPU IO ):
import base64
import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset
from torch.utils.data import IterableDataset
import torch.distributed as dist
class D(IterableDataset):
def __init__(self, path, number_of_sample, num_workers=4, training=False):
self.training = training
self.num_workers = num_workers
self.handles = [open(path, 'r') for _ in range(num_workers)]
self.LEN = number_of_sample
def __len__(self):
return self.LEN
def __iter__(self):
rank = dist.get_rank()
world_size = dist.get_world_size()
worker_info = torch.utils.data.get_worker_info()
mod = world_size
shift = rank
if worker_info:
mod *= worker_info.num_workers
shift = rank * worker_info.num_workers + worker_info.id
for i in range(self.LEN):
if (i + shift) % mod == 0:
yield self.__getitem__(i)
def __getline__(self, idx):
# 按照这个思路即可
return idx * 32
def __getitem__(self, idx, ):
worker_id = torch.utils.data.get_worker_info().id
# 读取数据
self.handles_feat[worker_id].seek(self.__getline__(idx))
data, label= self.handles_feat[d][worker_id].readline().strip().split('\t')
data = np.frombuffer(base64.b64decode(data), dtype=np.float32).ravel()
label = np.frombuffer(base64.b64decode(label), dtype=np.float32).ravel()
data = torch.FloatTensor(data)
label = torch.FloatTensor(label)
return {'data': data, 'label': label}
IterableDataset不支持shuffle
如果使用IterableDataset的话,则会遇到它并不支持shuffle的窘境。如果在内存中尝试对list(self.__len__)
进行shuffle的话,会生成一个数TB大小的index array,再加上训练的内存消耗,还没开始训练电脑已经重启了!
有没有方案能够不占用大量的内存,也可以对数据进行充分shuffle呢?如果我们不考虑丢失数据的话,可以直接在 数据size 的范围内随机生成index,但是这样的话,会丢失数据显然不是我们希望的结果。但是如果数据不进行充分shuffle的话,训练出来的模型效果会显著得变差。
方案一
在训练前独立进行数据的shuffle,例如通过 Map-Reduce 对于数据进行shuffle,这里的操作比较基础,就不给出实现代码了。
方案二
对数据进行分块shuffle,或者使用一个较大的buffer进行shuffle。但是分块如果太细的话,shuffle不够充分;分块太大的话,效率上又会显著低一些。而且性能上始终无法比得上完全随机shuffle。从实际跑出来的效果上来看,模型的性能确实下降了很多。
最终解决方案
这里的主要问题是,我们没办法把所有的index都放进内存进行排序,如果能存在某种函数,能实现在任意范围内建立一个映射关系,映射后的序列与原序列是一一对应,同时又“均匀”分布的话,是不是就能够实现shuffle数据的目标了?我们这里给出如下的解决方案。我们知道素数只能够被自己或者1整除,利用这个特性,我们可以建立一种一一映射的关系。
import numpy as np
N = 200
a = np.arange(N)
sep = 200
primenumber = 11
a_shuffle = (a%sep)*primenumber%sep + a//sep*sep
print('打乱后数组和元素组的不同元素个数:', len(set(a)-set((a%sep)*primenumber%sep + a//sep*sep)))
print(a_shuffle )
输出的结果为
打乱后数组和元素组的不同元素个数:0
array([ 0, 11, 22, 33, 44, 55, 66, 77, 88, 99, 110, 121, 132,
143, 154, 165, 176, 187, 198, 9, 20, 31, 42, 53, 64, 75,
86, 97, 108, 119, 130, 141, 152, 163, 174, 185, 196, 7, 18,
29, 40, 51, 62, 73, 84, 95, 106, 117, 128, 139, 150, 161,
172, 183, 194, 5, 16, 27, 38, 49, 60, 71, 82, 93, 104,
115, 126, 137, 148, 159, 170, 181, 192, 3, 14, 25, 36, 47,
58, 69, 80, 91, 102, 113, 124, 135, 146, 157, 168, 179, 190,
1, 12, 23, 34, 45, 56, 67, 78, 89, 100, 111, 122, 133,
144, 155, 166, 177, 188, 199, 10, 21, 32, 43, 54, 65, 76,
87, 98, 109, 120, 131, 142, 153, 164, 175, 186, 197, 8, 19,
30, 41, 52, 63, 74, 85, 96, 107, 118, 129, 140, 151, 162,
173, 184, 195, 6, 17, 28, 39, 50, 61, 72, 83, 94, 105,
116, 127, 138, 149, 160, 171, 182, 193, 4, 15, 26, 37, 48,
59, 70, 81, 92, 103, 114, 125, 136, 147, 158, 169, 180, 191,
2, 13, 24, 35, 46, 57, 68, 79, 90, 101, 112, 123, 134,
145, 156, 167, 178, 189])
我们这里得到的结果是对于原数据的一一映射,且会按照素数作为间隔采样。但是我们可能认为它不够“随机”,我们可以用到这里给出的sep超参数来加以优化,让它更加“随机”。这里整理成一个新的shuffle函数:
def shuffle_zhou(a, N, depth=1, seed=10086):
import random
random.seed(seed)
primes = [23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97,]
for d in range(depth):
sep = N//10**(depth-d-1)
primenumber = random.choice(primes)
a = (a%sep)*primenumber%sep + a//sep*sep
return a
N = 1000
a = np.arange(N)
import matplotlib.pyplot as plt
plt.plot(a, shuffle_zhou(a, N), '.')
输出如下,可以看出分布比较均匀,可以近似为随机分布。
总结
最终优化后的shuffleb版本IterableDataset如下,
import base64
import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset
from torch.utils.data import IterableDataset
import torch.distributed as dist
def shuffle_zhou(a, N, depth=3, seed=10086):
import random
random.seed(seed)
primes = [23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97,]
for d in range(depth):
sep = N//10**(depth-d-1)
primenumber = random.choice(primes)
a = (a%sep)*primenumber%sep + a//sep*sep
return a
class D(IterableDataset):
def __init__(self, path, number_of_sample, num_workers=4, training=False):
self.training = training
self.num_workers = num_workers
self.handles = [open(path, 'r') for _ in range(num_workers)]
self.LEN = number_of_sample
self.seed = 10086
def __len__(self):
return self.LEN
def __setseed__(self, seed=10086):
self.seed = seed
def __iter__(self):
rank = dist.get_rank()
world_size = dist.get_world_size()
worker_info = torch.utils.data.get_worker_info()
mod = world_size
shift = rank
if worker_info:
mod *= worker_info.num_workers
shift = rank * worker_info.num_workers + worker_info.id
for i in range(self.LEN):
i = shuffle_zhou(i, self.LEN, seed=self.seed)
if (i + shift) % mod == 0:
yield self.__getitem__(i)
def __getline__(self, idx):
# 按照这个思路即可
return idx * 32
def __getitem__(self, idx, ):
worker_id = torch.utils.data.get_worker_info().id
# 读取数据
self.handles_feat[worker_id].seek(self.__getline__(idx))
data, label= self.handles_feat[d][worker_id].readline().strip().split('\t')
data = np.frombuffer(base64.b64decode(data), dtype=np.float32).ravel()
label = np.frombuffer(base64.b64decode(label), dtype=np.float32).ravel()
data = torch.FloatTensor(data)
label = torch.FloatTensor(label)
return {'data': data, 'label': label}