「深度解析」AI训练之数据缓存

深度学习或者AI的出现，改变了我们以往的解决问题的编程方式，不再是代码上直观的表达。

举一个简单的例子，我们如何识别一个数字（图片）是数字9呢？非常直观的方法就是上面有一小圆圈，下面有一个竖线。但是如果写的倾斜了一些呢？如果上面的圈没有闭合呢？如果竖线弯曲了呢？感觉我们日常的程序判断（switch）无法收敛，我们只能用一种能够自我演进的方式来认识这个“看起来像9”的数字，而这也正是我们大脑的学习行为，我们第一个看到这个数字的时候，被告知这是9，那么图片就有了一个标签；下次再看到类似的，还是属于标签9，见多识广，最后见到一个也许写得更加不像的我们也能够识别出是9，这个过程正是由我们大脑中的上千亿的神经细胞长时间的学习结果。

人类大脑的真正运行方式，依旧是神秘所在，但从这个过程中我们发展出了神经网络算法，可以从已有的知识中进行学习。勤能补拙，既然算法不如人脑，就通过学习大量的资料来加快学习的进程。MNIST 数据有 60,000张手写数字图片，ImageNet 数据有接近1500万张图片，Youtube-8M的视频数据有数TB，Google的 Open Image dataset, 仅仅是在 Open Images Challenge中使用的数据集就达到了18TB。

AI中有三大核心：算法，算力，数据（存储）。算法自有成熟的框架，由数学科学家去解决；计算能力由CPU甚至GPU去解决。面对如此大量的数据，一台机器的内存、硬盘去承载基本不太可能，而对于CPU/GPU计算能力强悍的组件，频繁的去远端获取数据等待IO又是资源的浪费。有没有既能满足数据距离计算近、又能承载大量数据的方案呢？缓存是银弹！后面的主要篇幅从论文解析的角度来逐步阐述，论文来自Fast20 Quiver: An Informed Storage Cache for Deep Learning。

后续的讨论中，有个比较重要的概念，就是mini-batch, 如果没有实战的经历过，不是很容易理解这个概念。

深度学习的优化算法，本质就是梯度下降。每次的参数更新有两种方式。

第一种，遍历全部数据集计算一次损失函数，然后计算函数对各个参数的梯度，更新梯度。这种方法每更新一次参数都要把数据集里的所有的样本数据遍历一遍，计算量开销大，计算速度慢，不支持在线学习，这称为Batch gradient descent（BGD），批梯度下降。

另一种，每训练一个数据就算一下损失函数，然后求梯度更新参数，这个称为随机梯度下降，stochastic gradient descent。这个方法速度比较快，但是收敛性能不太好，可能在最优点附近波动，达不到最优点。两次参数的更新也有可能互相抵消掉，造成目标函数震荡的比较剧烈。

为了克服两种方法的缺点，现在一般采用的是一种折中手段，mini-batch gradient decent，小批的梯度下降，这种方法把数据分为若干个批，按批来更新参数，这样，一个批中的一组数据共同决定了本次梯度的方向，下降起来就不容易跑偏，减少了随机性。

用一个示意图表示如下：

蓝色：为 batch 梯度下降，即 mini batch size ＝ m，紫色：为 stochastic 梯度下降，即 mini batch size ＝ 1，绿色：为 mini batch 梯度下降，即 1 < mini batch size < m。

以下图为例，执行了3轮训练（epoch）,每轮里面定义mini-bach size=5, 其中数据集为1-20个数字，我们看到通过torch.DataLoader, 每次获得了5个数据(batch x)。

01 深度学习训练的基本知识

深度学习训练任务(Deep Learning Training DLT)会将训练数据作为输入，从千丝万缕的线索中通过学习并得到一个输出模型来代表训练数据。

为了实现训练，DLT会使用一个较小的随机样本（mini-batch,通常是32到512个），并利用SGD来慢慢的学习各种参数进而提高准确率。

训练数据：通常我们可以认为是一个列表，列表中的每一个元素都是一个二元组<input,label>, input可能是一张图片或者一段语音，而label则代表着input的语义，而这也正是深度学习网络所需要学习的并能够正确区分input的目标。例如ImageNet的全部数据集大概有150万张图片，每张图片在200KB左右。

为了能够以随机方式访问训练数据，DLT框架会使用索引序列来遍历数据。假设训练数据有100万个文件，那么会维护一个包含每一个文件索引的列表，并对它进行随机的排列，随后根据mini-batch的数据量向后端存储获得数据，当全部的数据都完整遍历训练一次，一个epoch完成。对于下一个epoch, 再次对索引进行随机排列，重复上面的过程。一个典型的DLT任务会运行很多轮训练，例如50-200。

数据转换：从存储获得原始数据会被DLT框架进行转换，例如彩色图片变成黑白图片，同时将图片转换为像素数矩阵等等。当然这部分工作通常由CPU来完成。

多任务：因为DLT任务是一个试错的过程，所以实际运行过程中，用户总是会使用不同的参数来同时运行不同的任务，所有的这些任务都会访问相同的完整数据集，不同的就是以不同的随机顺序来进行访问。

02 深度学习的IO特点

我们从DLT任务I/O访问的角度看来列举一下它的主要特点：

可共享性：在DLT训练任务中，无论是一个训练任务自身，还是多个训练任务之间，都存在很大程度的I/O重叠性。在一个任务内，它会针对同一个数据集进行多次的遍历（例如多个epoch），所以如果能够在第一个epoch的时候就对数据进行缓存，会大幅提升后续训练的效率。更重要的是，这种共享性甚至可以扩展到多任务之间，例如针对同一份训练数据集，配置不同的参数，利用同一个训练模型运行多个不同的任务。这些任务可能运行在不同的机器上，但是访问的都是相同的底层数据集。

随机访问：由于数据的可共享性，这使得DLT具有非常的缓存友好性，但只有在全部数据能够被完整缓存的情况下才有效果，否则，DLT随机访问数据的方式又使得部分数据缓存很容易被穿透。例如只能够缓存20%的数据，那么这些数据马上就会被后续的随机访问刷掉。

部分数据缓存对于DLT来说很重要，因为训练数据通常已经足够大，并且会越来越多，例如前文提到过即使只是ImageNet这样的百万级规模的数据集，总体也已经达到了数TB的大小。

可替换性：从I/O的角度来说，一个训练任务(epoch)主要关注以下两点即可：a)每一个训练数据必须且仅被访问一次；b)而对于每次的mini-batch,必须是随机的序列。有趣的是，一组数据的精确顺序并不会对训练任务的准确或者精确性产生影响，这就意味着I/O是可以被替换的。对于特定若干文件的访问，DTL任务可以替换为一组其他的随机的且没有被访问过的数据。从缓存的角度来说，这是一个独特的特性来提升缓存的命中率，即使缓存只能承载20%的数据，也可以在访问一个不存在于缓存中的数据，通过替换的方式返回一个存在的内容，同时并没有破坏随机以及唯一性的训练要求。

可预测性：因为每一个mini-batch的运行时间，可以事先获得，这样就可以用于评估一个训练任务对I/O性能的敏感性，进而可以进行策略调整以能够使那么I/O敏感的任务从缓存获益。

03 缓存的设计

总结起来深度学习的特点：

1. 需要的数据量大
2. 多台机器多个训练并行
3. 每个训练要运行多次
4. 在每次训练中，所有的数据需要遍历一遍
5. 针对不同的训练参数，以及在不同的机器上运行的训练任务，数据集相对保持固定

针对以上的特点，当我们考虑缓存的时候，不禁会有如下的疑问：缓存毕竟容量有限，穿透如何处理？缓存的过期置换策略是如何的？当不同的用户访问不同的数据，安全性如何保证？等等。

Quiver、分布式缓存，通过与DLT框架深度整合，缓存客户端集成到训练任务的IO过程中，进而为缓存服务端提供更多的策略信息。

系统结构

以公有云虚拟机环境举例，每一个GPU VM带有一个本地SSD硬盘，每一个DLT job会运行在自己的容器内，这样即使是多用户运行，也是在一个隔离的环境内。

每个用户的数据存储在各自账号的云存储内，这样保证了隐私以及访问权限。通过分布式缓存，即使训练任务由于调度等原因在各个宿主之间切换，缓存数据依旧是能够提高训练效率。

数据安全

Quiver的设计是一个共享式的分布缓存，无论是不同的任务，还是不同的用户之间，在共享的模式下如何保证数据的安全就是一个重要因素。Quiver保证了用户只能看到他有权限访问的数据，但这样似乎又与缓存的重用产生了冲突。如果针对某一个数据集，例如ImageNet, 两个不同用户分别在各自的存储账号内各自保存了一份，那么逻辑上来讲，缓存要分别为每个用户各自缓存一份。这将导致缓存的效率降低，Quiver通过内容寻址(content-addressed)的方式来解决重用与隔离的问题。

内容寻址缓存

对于缓存，基本的行为就是通过一个<key, value>的映射关系，在我们通过key查询时，能够快速的返回所对应的value。在Quiver中，缓存并不是利用文件名以及偏移量在作为缓存的关键字，而是利用缓存内容的hashes。缓存内容的粒度是由具体的DLT任务决定的，可能是一张图片，无论是对它的插入还是寻址，都以它的hash(例如SHA1)来唯一定位。用hash来定位的好处是对于一个相同内容的文件，不管它来自于何处以及文件名是否相同，在缓存中都仅需要保留一份即可，这样也就能够达到即使在不同的用户之间也能够共享目的。

同时为了保证数据的隔离性，Quiver利用摘要索引来访问训练数据，对于每一份数据，摘要索引将包含<content_hash: file_location>， 因此，在多用户拥有相同内容的数据集时，因为数据是存在在各自的存储系统内，每个用户将拥有不同的file_location，但是所有的content_hash是相同的。

缓存服务器

利用本地SSD作为介质的KV存储，通过一致性Hash的方式将key space分布在多个缓存服务器上。

缓存管理(Cache Manager)

由于Quiver是分布式缓存，那么针对所有的缓存服务器，缓存的插入、清理需要一个协调者Cache manager。

Cache manager同时会评估每一个计算任务从缓存的受益情况，主要通过让缓存服务器针对训练任务所需要的若干mini-batch数据做cache misses, 然后与其他的缓存命中的训练人耗时机型对比，进而对缓存进行优先级调整。

缓存客户端

缓存客户端作为训练任务的一部分，通过干预DLT框架，例如PyTorch等的接口层来访问训练数据。在PyTorch中，DataSet会用来遍历所有的训练数据，并且内部维护一个随机的文件索引列表，其中Next的接口就可以用来获得下一个mini-batch数据。Quiver通过调整这个接口，利用一个摘要文件，当上层访问一组文件时，它会先对缓存进行数据的访问。

客户端会将训练任务的一些信息反馈给Cache Manager,例如每一个mini-batch的训练时间，Cache Manager可以据此来优化缓存的策略。

替换命中率

在常规的缓存中，如果一个mini-batch包含了512个文件，那么Dataset会提供512个文件索引用来从后端存储获得文件内容，如果这其中只有部分缓存命中，那么将依然存在远程的I/O操作。在Quiver中，会从Cache中加载更多的（例如10倍的mini-batch数量）数据，而只要其中有512个数据能够被命中，那么就返回给上层训练任务，这样训练任务就不会被Cache miss阻塞。同时Quiver会标记Cache miss的数据为pending状态，周而复始，直到数据被遍历了一遍，这时将重头来过仅仅去关注之前pending的数据。

假设目前只有10%的数据在缓存中，为了简单起见，我们可以认为就是连续的原始数据的10%， 因为DLT任务会随机的查找数据，所以每一个长度为k的mini-batch序列，缓存的命中率应该为k/10, 因此如果我们查找一个长度为10*k的序列，那么正好能够命中获得mini-batch所需要的数据。当下一轮查找pending数据的时候，另外的10%的数据可能已经在缓存中了，这也意味着能够1/9的命中率。需要注意的是，在多个任务的训练中，这依旧适用，因此多个训练任务尽管每个都访问随机的训练数据，从整体来看，他们可以做到以全缓存命中的方式来运行。

训练准确性

由于上述的I/O可替换性，我们有理由怀疑最终训练结果的准确性。这里借用原文的数据来说明。

04 缓存的管理

在之前的描述中，当只有部分数据被缓存时，Quiver会在一个epoch的训练过程中，再次遍历文件索引。为了能在这后续的遍历中获得更好的命中率，另一部分数据必须被pre-fetch到缓存中。

Quiver通过缓存整个数据集的2个chunks来解决这个问题。首先数据集的摘要索引文件会被分成固定数据的chunks,例如每个chunk包含10%的数据，同时每个chunk代表着striped partition。例如我们定义数据集中连续的10%为一个partition, 每个partition被分成10个stripe units. 这样每个chunk将包含所有的partition中的一个unit。这样当训练任务操作第一个chunk的过程中，第二个chunk将被加载到缓存中，所以当部分训练任务完成第一次遍历开始第二次的时候，数据已经在缓存内，训练任务以递进的方式运行。

这里面潜在的一个问题就是什么时候将第一个chunk置换出去。如果置换太快，部分任务还没有完成将导致缓存失效，如果保留太长时间，那么第三个chunk将无法加载进来。在Quiver中，当第二个chunk被加载到缓存后，第一个chunk会被标记为可以清除，同时新的任务可以从第二个chunk中获得命中的数据。原有的任务依旧利用第一个chunk来运行，当所有的任务都已经遍历了第一个chunk数据，这部分数据才会真的从缓存中清除，同时第三部分数据将开始加载。

在上述的过程中，如果某一个训练任务相比于其他的要慢很多，那么将导致前一个chunk迟迟不能释放，通常来说，在同一个训练模型的多个任务中，每个任务的训练时间基本是相同的，但无法避免在多个不同的训练模型训练同一个数据集的场景。不过如果一个任务明显的耗时很长，那么将意味着每一个mini-batch在GPU上的训练时间都很长，也就是它对I/O的性能没那么敏感，所以缓存的不命中并不会影响多少这个训练的效率，因此 Quiver会设定一个临界值来强制第一个chunk失效并清除。

05 缓存效果

论文作者通过如下的配置环境来进行效果的对比，从实际数据来看，训练性能确实有较大的提高。

Timeline of Mini-batches

吞吐提升

06 结论

深度学习场景中，更多的注意力放在了提高计算以及网络的性能上，而对于存储，则是利用现有的方案来解决，例如提前手动将数据加载到离GPU较近的SSD上。论文作者通过Quiver，提供了自动化的手段来消除存储的瓶颈。当然无法避免对训练框架的侵入性，但也正是因为如此，Quiver才能感知到训练I/O的特点，进而达到即使缓存只能承载部分数据也可以大幅调高缓存利用率的效果。