MLX对比MPS对比CUDA：基准测试

苹果新的ML框架MLX的首个基准测试

如果你是Mac的用户和深度学习爱好者，你可能曾经希望你的Mac能够处理那些庞大的模型，对吧？那么，你知道吗？苹果刚刚发布了MLX，这是一个能够高效运行苹果芯片上的机器学习模型的框架。

在PyTorch 1.12中最近引入的MPS后端已经是一个大胆的步骤，但是随着MLX的宣布，似乎苹果想要在开源深度学习领域迈出重要的一步。

在本文中，我们将对这些新的方法进行基准测试，与传统的CPU后端和两个支持CUDA的GPU进行对比。通过这样做，我们的目标是揭示这些新的适用于Mac的方法在2024年在深度学习实验中可以起到多大的作用。

作为一个以GNN为导向的研究者，我将重点考察图卷积网络（GCN）模型的基准测试。但是由于这个模型主要由线性层组成，我们的发现对于那些不特别在GNN领域的人也可能具有启发意义。

构建环境

为了为MLX构建一个环境，我们需要指定是使用i386还是arm架构。使用conda，可以通过以下方式实现：

CONDA_SUBDIR=osx-arm64 conda create -n mlx python=3.10 numpy -c conda-forgeconda activate mlx

要检查你的环境是否真的使用了arm，以下命令的输出应该是arm，而不是i386：

python -c "import platform; print(platform.processor())"

现在只需使用pip安装MLX，你就可以开始探索了：

pip install mlx

GCN实现

GCN模型是一种图神经网络（GNN）模型，它使用邻接矩阵（表示图的结构）和节点特征进行工作。它通过从相邻节点收集信息来计算节点嵌入。具体地说，每个节点会得到其邻居特征的平均值。这个平均值是通过将节点特征与归一化的邻接矩阵相乘来计算的，同时考虑了节点的度。为了学习这个过程，首先将特征投影到一个嵌入空间中，通过线性层实现。