DeepSpeed Megatron和pytorch原生的分布式区别

转载

mob64ca140f67e3 2024-10-29 21:55:57

文章标签 人工智能 git python 2d ide 文章分类 PyTorch 人工智能

本文介绍了最近更新的 DiracNet 实现项目，该项目实现了不带跳过连接的超深层网络，并且是对应论文的官方实现。机器之心简要介绍了该项目和论文。

如 ResNet 这样采取了跳过连接（skip-connections）的网络在图像识别基准上实现了非常优秀的性能，但这种网络并体会不到更深层级所带来的优势。因此我们可能会比较感兴趣如何学习非常深的表征，并挖掘深层网络所带来的优势。我们提出了一个简单的权重参数化（weight parameterization）方法，该方法将提升不带跳过连接的深层网络性能，且允许我们训练数百层的深度网络。我们提出来的 DiracNet 非常接近于宽 ResNet 的准确度，但也需要更多的参数。此外，我们发现只有 34 层的 DiracNet 在性能上要优于 ResNet-1000。我们提出的 Dirac 权重参数化在推断时能折叠进一个卷积核，因此它可以轻易地解释为类 VGG 的网络。

下图展示了随层级和宽度的变化，CIFAR-10 准确度的变化。

DeepSpeed Megatron和pytorch原生的分布式区别_ide

简介

简而言之，Dirac 参数化就是卷积核权重的求和并通过 Dirac delta 函数缩放：

conv2d(x, alpha * delta + W)

以下是简化后的 PyTorch 伪代码，其展示了我们用于训练 plain DiracNets（带有权重归一化）的函数：

def dirac_conv2d(input, W, alpha, beta)
    return F.conv2d(input, alpha * dirac(W) + beta * normalize(W))

其中 alpha 和 beta 为缩放标量，normalize 为每一个特征图上的 L2 归一化。此外，我们还使用了 NCReLU（负的 CReLU）非线性激活函数：

def ncrelu(x):
    return torch.cat([x.clamp(min=0), x.clamp(max=0)], dim=1)

代码

在本项目中，diracconv.py 文件为模块化的 DiracConv 定义、test.py 为单位测试文件、 diracnet-export.ipynb 文件为 ImageNet 预训练模型、 diracnet.py 文件为函数化的模型定义，最后的 train.py 为 ImageNet 训练代码。

依赖库的安装

首先实现本项目的模型需要已安装 PyTorch，但这里并不会详细解释如何安装这个模块。随后我们可以使用 pip 安装 torchnet：

pip install git+https://github.com/pytorch/tnt.git@master

使用 Python 捆绑包安装 OpenCV，即执行命令 conda install -c menpo opencv3（前提已安装 Anaconda），然后使用 OpenCV 转换 torchvision：

pip install git+https://github.com/szagoruyko/vision.git@opencv

最后，安装其它的 Python 包：

pip install -r requirements.txt

在 CIFAR 上训练 DiracNet-34-2 请运行以下命令：

python train.py --save ./logs/diracnets_$RANDOM$RANDOM --depth 34 --width 2

在 ImageNet 上训练 DiracNet-18 请运行以下命令：

python train.py --dataroot ~/ILSVRC2012/ --dataset ImageNet --depth 18 --save ./logs/diracnet_$RANDOM$RANDOM \
                --batchSize 256 --epoch_step [30,60,90] --epochs 100 --weightDecay 0.0001 --lr_decay_ratio 0.1

nn.Module 代码

我们还提供了 DiracConv1d、DiracConv2d 和 DiracConv3d 等 API，它们使用的方法就像 nn.Conv1d、nn.Conv2d 和 nn.Conv3d，只不过我们提供的 API 使用的是 Dirac-parametrization。我们的训练代码并没有使用这些模块，但它十分有助于各位读者快速地实现 Dirac 层级。

预训练模型

我们将批量归一化和 Dirac 参数层级放入了 F.conv2d 的 weight 和 bias 张量中，因此我们的模型就像 VGG 或 AlexNet 那样简单方便地调用，它们都只调用了最基本的非线性和 conv2d 模块。diracnets.ipynb 文件中可以查看模型定义的函数与模块。

下面我们提供了 DiracNet-18-0.75 序列模型的打印输出结果，可作为读者的参考：

Sequential (
  (conv): Conv2d(3, 48, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3))
  (max_pool0): MaxPool2d (size=(3, 3), stride=(2, 2), padding=(1, 1), dilation=(1, 1))
  (group0.block0.ncrelu): NCReLU()
  (group0.block0.conv): Conv2d(96, 48, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (group0.block1.ncrelu): NCReLU()
  (group0.block1.conv): Conv2d(96, 48, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (group0.block2.ncrelu): NCReLU()
  (group0.block2.conv): Conv2d(96, 48, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (group0.block3.ncrelu): NCReLU()
  (group0.block3.conv): Conv2d(96, 48, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (max_pool1): MaxPool2d (size=(2, 2), stride=(2, 2), dilation=(1, 1))
  (group1.block0.ncrelu): NCReLU()
  (group1.block0.conv): Conv2d(96, 96, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (group1.block1.ncrelu): NCReLU()
  (group1.block1.conv): Conv2d(192, 96, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (group1.block2.ncrelu): NCReLU()
  (group1.block2.conv): Conv2d(192, 96, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (group1.block3.ncrelu): NCReLU()
  (group1.block3.conv): Conv2d(192, 96, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (max_pool2): MaxPool2d (size=(2, 2), stride=(2, 2), dilation=(1, 1))
  (group2.block0.ncrelu): NCReLU()
  (group2.block0.conv): Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (group2.block1.ncrelu): NCReLU()
  (group2.block1.conv): Conv2d(384, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (group2.block2.ncrelu): NCReLU()
  (group2.block2.conv): Conv2d(384, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (group2.block3.ncrelu): NCReLU()
  (group2.block3.conv): Conv2d(384, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (max_pool3): MaxPool2d (size=(2, 2), stride=(2, 2), dilation=(1, 1))
  (group3.block0.ncrelu): NCReLU()
  (group3.block0.conv): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (group3.block1.ncrelu): NCReLU()
  (group3.block1.conv): Conv2d(768, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (group3.block2.ncrelu): NCReLU()
  (group3.block2.conv): Conv2d(768, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (group3.block3.ncrelu): NCReLU()
  (group3.block3.conv): Conv2d(768, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (relu): ReLU ()
  (avg_pool): AvgPool2d ()
  (view): Flatten()
  (fc): Linear (384 -> 1000)
)

论文：DiracNets: Training Very Deep Neural Networks Without Skip-Connections

DeepSpeed Megatron和pytorch原生的分布式区别_ide_02

论文地址：https://arxiv.org/abs/1706.00388

我们为神经网络提出了一种新的权重参数化（weight parameterization）方法，这种方法可称之为 Dirac 参数化，它适用于广泛的网络架构。此外，通过使用上面的参数化方法，我们提出了一种新的 plain VGG-like 和 ResNet-like 架构。该架构没有使用任何跳过连接（skip-connections），我们称之为 DiracNet。这些网络可以训练数百层，且在 34 层的时候就超过了 1001 层的 ResNet 性能，且十分接近宽 ResNet（WRN）的准确度。我们应该注意训练 DiracNets 是一种端到端的方法，并不需要任何层级的预训练方法。我们相信该研究是迈向更简单和高效深度神经网络的重要方向。

DeepSpeed Megatron和pytorch原生的分布式区别_python_03