今天,分享一篇训练速度最高100倍提升!基于PyTorch实现的可微逻辑门网络开源,希望以下训练速度最高100倍提升!基于PyTorch实现的可微逻辑门网络开源的内容对您有用。
编辑:袁铭怿
本文通过逻辑门组合学习来探索面向机器学习任务的逻辑门网络。这些网络由 AND 和 XOR 等逻辑门组成,为了实现有效训练,本文提出可微逻辑门网络,一种结合了实值逻辑和网络连续参数化松弛的架构。
随着神经网络的成功应用,各项研究和机构也一直在致力于实现快速且高效的计算,特别是在推理时。对此,各种技术应运而生,包括降低计算精度,二进制和稀疏神经网络。本文中,来自斯坦福大学、萨尔茨堡大学等机构的研究者希望训练出一种不同的、在计算机领域被广泛应用的体系架构:逻辑门网络(logic (gate) networks)。
训练像逻辑门这样的离散组件网络所面临的问题是,它们是不可微的,因此,一般而言,不能通过梯度下降等标准方法进行优化。研究者提出了一种方法是无梯度优化方法,如演化训练(evolutionary training),它适用于小型模型,但不适用于大型模型。
在这项工作中,该研究探索了用于机器学习任务的逻辑门网络。这些网络由「AND」和「XOR」等逻辑门电路组成,可以快速执行任务。逻辑门网络的难点在于它通常是不可微的,不允许用梯度下降进行训练。因此,可微逻辑门网络的出现是为了进行有效的训练。由此产生的离散逻辑门网络实现了快速的推理速度,例如,在单个 CPU 核上每秒处理超过一百万张 MNIST 图像。这篇论文入选 NeurIPS 2022 。
论文地址:https://arxiv.org/pdf/2210.08277.pdf
项目地址:https://github.com/Felix-Petersen/difflogic
纽约大学计算机科学教授 Alfredo Canziani 表示:由逻辑门(如 AND 和 XOR)组成的可学习组合网络,允许非常快速的执行任务及硬件实现。
距离论文公开才过去三个月,论文作者 Felix Petersen 表示该研究的官方实现已经公布,他们发布了 difflogic 项目,这是一个基于 pytorch 实现的可微逻辑门网络库。不仅如此,经过优化现在的训练速度比最初的速度快 50-100 倍,因为该研究提供了高度优化的 CUDA 内核。
项目介绍
difflogic 是一个基于 Python 3.6 + 和 PyTorch 1.9.0 + 的库,基于逻辑门网络进行训练和推理。该库安装代码如下:
pip install difflogic
需要注意的是,使用 difflogic,还需要 CUDA、CUDA 工具包(用于编译)以及 torch>=1.9.0(匹配 CUDA 版本)。
下面给出了 MNIST 数据集的可微逻辑网络模型的定义示例:
from difflogic import LogicLayer, GroupSum
import torch
model = torch.nn.Sequential(
torch.nn.Flatten(),
LogicLayer(784, 16_000),
LogicLayer(16_000, 16_000),
LogicLayer(16_000, 16_000),
LogicLayer(16_000, 16_000),
LogicLayer(16_000, 16_000),
GroupSum(k=10, tau=30)
)
该模型接收了 784 维的输入,并返回 k=10,对应于 MNIST 的 10 个类。该模型可以用 torch.nn.CrossEntropyLoss 进行训练,类似于在 PyTorch 中训练其他神经网络模型的方式。值得注意的是,Adam 优化器(torch. optimt .Adam)可用于训练,推荐的默认学习率是 0.01 而不是 0.001。最后,同样需要注意的是,与传统的 MLP 神经网络相比,逻辑门网络每层神经元的数量要高得多,因为后者非常稀疏。
为了深入了解这些模块的细节,下面是一些更详细的示例:
layer = LogicLayer(
in_dim=784, # number of inputs
out_dim=16_000,# number of outputs
device='cuda', # the device (cuda / cpu)
implementation='cuda', # the implementation to be used (native cuda / vanilla pytorch)
connections='random',# the method for the random initialization of the connections
grad_factor=1.1, # for deep models (>6 layers), the grad_factor should be increased (e.g., 2) to avoid vanishing gradients
)
模型推理
在训练期间,模型应该保持在 PyTorch 训练模式,即.train (),这种模式使得模型保持可微。现在有两种模式可以进行快速推理:
第一种选择是使用 PackBitsTensor。PackBitsTensors 允许在 GPU 上高效动态的执行训练好的逻辑门网络。
第二种选择是使用 CompiledLogicNet。CompiledLogicNet 允许在 CPU 上高效地执行固定训练的逻辑门网络。
下面给出一些实验示例,它们包含在 experiments 目录中。main.py 用于执行,main_baseline.py 是包含规则的神经网络基线。
MNIST
python experiments/main.py -bs 100 -t 10 --dataset mnist20x20 -ni 200_000 -ef 1_000 -k 8_000 -l 6 --compile_model
python experiments/main.py -bs 100 -t 30 --dataset mnist-ni 200_000 -ef 1_000 -k 64_000 -l 6 --compile_model# Baselines:
python experiments/main_baseline.py -bs 100 --dataset mnist -ni 200_000 -ef 1_000 -k 128 -l 3
python experiments/main_baseline.py -bs 100 --dataset mnist -ni 200_000 -ef 1_000 -k 2048 -l 7
CIFAR-10
python experiments/main.py -bs 100 -t 100 --dataset cifar-10-3-thresholds -ni 200_000 -ef 1_000 -k 12_000 -l 4 --compile_model
python experiments/main.py -bs 100 -t 100 --dataset cifar-10-3-thresholds -ni 200_000 -ef 1_000 -k128_000 -l 4 --compile_model
python experiments/main.py -bs 100 -t 100 --dataset cifar-10-31-thresholds -ni 200_000 -ef 1_000 -k256_000 -l 5
python experiments/main.py -bs 100 -t 100 --dataset cifar-10-31-thresholds -ni 200_000 -ef 1_000 -k512_000 -l 5
python experiments/main.py -bs 100 -t 100 --dataset cifar-10-31-thresholds -ni 200_000 -ef 1_000 -k 1_024_000 -
想要了解更多内容,请参考原项目。