Skip to content
返回

一种针对主干网络的缓存机制的简单实现

发布于

在情绪评估模型的训练和推理过程中,编码器网络通常占据了模型计算资源的主要部分。 特别是在知识蒸馏这类任务中,教师网络往往是完全冻结的,这意味着对于相同的输入,教师网络总是产生相同的输出特征。 这些冗余计算不仅浪费了计算资源,还显著延长了模型训练时间,降低了研究和开发效率。

特征缓存系统包括内存缓存、设备缓存和磁盘缓存三个层次,可以有效地存储和复用已计算的特征张量。 该系统使用哈希函数为每个输入数据生成唯一的缓存键,避免重复计算。

缓存系统工作的基本流程

图 1 缓存系统工作的基本流程

图 1 展示了缓存系统工作的基本流程。 缓存系统的核心思想是避免在相同输入数据上重复执行计算成本高昂的特征提取操作,特别是当编码器被冻结时,对相同输入的特征输出应保持不变。

在每个训练与验证阶段,模型会接收原始数据,这些数据首先经过预处理,然后送入编码器进行特征提取。 在第一轮训练和验证阶段,由于缓存尚未建立,所有原始特征需要完整计算。计算完成后,系统会为每个输入数据生成唯一的缓存键(基于输入数据和模型参数的哈希值),并将对应的特征向量结果存入缓存系统。 在缓存空间足够的情况下,这些计算结果会被储存在三级缓存架构中,优先使用GPU显存,其次是主机内存,最后是磁盘存储。 在第二轮及后续训练和验证阶段,系统首先检查输入数据的缓存键是否存在于缓存中。 如果没有出现缓存空间被占满导致的缓存淘汰情况,这些缓存将会全部命中,从而完全跳过数据预处理和编码器模型推理的过程,直接返回缓存的特征向量。

这种机制在大型模型训练中可显著减少计算时间,尤其对于计算密集型的视觉和语音编码器,加速效果更为明显。 当缓存空间不足时,系统会基于最近最少使用(Least Recently Used,LRU)策略动态管理缓存内容,优先保留近期频繁访问的数据特征,淘汰长时间未使用的缓存项。 这种自适应缓存管理机制能够在有限的存储资源约束下最大化缓存效率。 此外,框架中会自动检测是否符合缓存条件,包括两个关键判断:一是编码器模型是否被冻结(确保相同输入产生相同输出的前提),二是缓存系统标志是否开启(允许用户灵活控制缓存功能)。 只有当这两个条件同时满足时,系统才会激活缓存系统,此时所有的前向传播请求都会被缓存管理器截获并进行缓存查找,从而实现特征计算的高效复用。 这种设计既保证了计算正确性,又最大化提升了训练效率。

缓存管理器

图 2 缓存管理器运行机制

本章的缓存系统是通过拦截前向传播过程并调用缓存管理器实现的, 图 2 展示了缓存管理器的运行机制和实现原理, 缓存系统采用了三层架构设计:GPU显存作为一级缓存,储存最近使用的特征张量,提供最快的访问速度;内存缓存作为第二级缓存,存储容量大于一级缓存,但访问速度较慢,需要额外的CPU到GPU数据传输;磁盘缓存提供最大的存储容量,但访问速度最慢。 当设备缓存未命中时,系统会检查内存缓存,若内存缓存也未命中,则尝试从磁盘加载缓存数据。 在计算哈希值时,缓存管理器将使用 SafeTensors¹ 对编码器模型对象进行序列化,转换为二进制数据, 然后将其与原始数据的元信息进行拼接,然后通过 SHA-256 算法对其进行处理生成哈希值。

在写入缓存时,缓存管理器需要接受三个参数:编码器模型对象、原始数据的元信息和要缓存的特征向量结果, 缓存管理器可以通过分别对编码器模型对象和原始数据的元信息计算哈希值并拼接来生成唯一的缓存哈希值,并将其作为缓存键, 将特征向量结果作为缓存值存储在缓存中,缓存会按照显存、内存、磁盘的顺序进行存储,其中磁盘的缓存值需要使用 SafeTensors¹ 对向量进行序列化后再储存。 如果缓存空间不足,缓存管理器会根据最近使用(Least Recently Used,LRU)策略清除最久未使用的缓存项,以腾出空间存储新的缓存数据。

在读取缓存时,缓存管理器需要接受两个参数:编码器模型对象和原始数据的元信息, 缓存管理器可以通过分别对编码器模型对象和原始数据的元信息计算哈希值并拼接得到唯一的缓存哈希值, 缓存管理器会这个哈希值查找缓存键,并按层级顺序检查显存、内存和磁盘缓存,直到找到匹配的缓存键为止。 如果缓存命中,缓存管理器会将缓存值返回给调用者,并跳过编码器模型的前向传播过程。 如果缓存未命中,缓存管理器会将原始数据传递给编码器模型进行前向传播,并将计算得到的特征向量结果写入缓存。

延迟加载机制

为了进一步提高缓存效率, 缓存管理器使用了延迟加载机制, 即仅在需要时才将特征从磁盘加载到内存或设备内存,以避免不必要的数据传输。 通过多层次特征缓存系统,可以显著减少深度学习模型训练中的计算冗余,提高训练效率,特别是在涉及大型冻结网络的场景中。 这一技术不仅加速了模型训练过程,还优化了计算资源的利用,为深度学习模型的高效训练和推理提供了新的解决方案。

源代码: emotion-recognition


建议修改

下一篇
将 qcow2 镜像导入 incus