中文

conv_einsum:卷积张量神经网络中多线性运算表示与快速评估的框架

机器学习 2024-01-09 v1 计算机视觉与模式识别

摘要

现代 ConvNets 在各种视觉和图像分类任务中持续取得最先进(SOTA)的结果,但代价是参数量的增加。在不牺牲太多表达能力的情况下压缩网络的一种策略是将其重塑为张量神经网络(TNN),即对其层进行高阶张量化,随后进行分解(如 CP 分解),从而将权重剥离为其关键基分量。通过 TNN 的前向传播可以表示为一系列多线性运算(MLOs),其中评估路径会极大地影响所产生的浮点运算次数(FLOPs)。虽然像流行的 einsum 这样的函数可以评估收缩等简单的 MLOs,但现有的实现无法处理多路卷积,导致缺乏对通过张量化卷积层的最优评估路径如何提高训练速度的评估。在本文中,我们开发了一个统一框架,将张量卷积层表示为类似 einsum 的字符串,并提出了元算法 conv_einsum,能够以最小化 FLOPs 的方式评估这些字符串。使用我们的开源实现,在广泛的模型、张量分解和多样化任务上进行的综合实验表明,conv_einsum 显著提高了卷积 TNN 的计算和内存效率。

关键词

引用

@article{arxiv.2401.03384,
  title  = {conv_einsum: A Framework for Representation and Fast Evaluation of Multilinear Operations in Convolutional Tensorial Neural Networks},
  author = {Tahseen Rabbani and Jiahao Su and Xiaoyu Liu and David Chan and Geoffrey Sangston and Furong Huang},
  journal= {arXiv preprint arXiv:2401.03384},
  year   = {2024}
}