Pytorch DenseNet训练自己的数据集(续)

千穹凌帝

已于 2024-06-16 12:19:36 修改

阅读量976

点赞数 15

文章标签： pytorch 人工智能 python

于 2024-06-16 11:55:02 首次发布

本文链接：https://blog.csdn.net/m0_59092412/article/details/139710220

版权

《画堂春》

灵珠有泪自千行等闲芳草斜阳离人过客暗凄凉偷羡鸳鸯

伤心脉脉难诉风剪寸寸柔肠神仙人鬼两茫茫情短恨长

之前写了一篇，使用pytorch实现DenseNet，完成完整的代码框架，从建立数据集、设置参数、训练网络到推理测试。链接如下：

https://blog.csdn.net/m0_59092412/article/details/139717977?csdn_share_tail=%7B%22type%22%3A%22blog%22%2C%22rType%22%3A%22article%22%2C%22rId%22%3A%22139717977%22%2C%22source%22%3A%22m0_59092412%22%7D

使用上没有问题，就是显卡消耗太大，2块RTX TITAN 24G只能跑最小的net121，batch只能到4，训练速度也慢，这还是用了DP并行加速后的结果，唉，一声长叹，丹炉不行，奈若何，换做药尘也无力。

手握泰坦，以为天下事无可不为，一丹出炉，方知天下人力有不足 --by 王权霸业

小门小派的悲哀，算了，既然丹炉没条件，只能靠炼丹术找出路。因此，在上一版的基础上，做些优化改进，设法提升训练速度，降低显存消耗。

主要策略：

1.使用Nvidia的APEX DDP库，利用distributed data parallel并行加速，混合精度以及同步BN；

2.将网络中的标准卷积换成深度可分离卷积；

3.使用凯明权值初始化方法；

4.增加SWISH和PReLU激活函数；

代码放在GitHub上，路径如下：

GitHub - tklk610/The-second-edition-of-DenseNet-with-Pytorch

APEX 是 NVIDIA 开源的用于混合精度训练和分布式训练库。Apex 对混合精度训练的过程进行了封装，改两三行配置就可以进行混合精度的训练，从而大幅度降低显存占用，节约运算时间。此外，Apex 也提供了对分布式训练的封装，针对 NVIDIA 的 NCCL 通信库进行了优化。

可参考如下：

英伟达（NVIDIA）训练深度学习模型神器APEX使用指南_查看apex版本 nivida-CSDN博客

从NVIDIA的github安装apex库：

git clone https://github.com/NVIDIA/apex
cd apex

# 同时安装C++扩展
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
 
# Apex 同样支持 Python-only build (required with Pytorch 0.4) via
 
pip install -v --no-cache-dir ./

APEX中包含了混合精度，混合精度训练，指代的是单精度 float和半精度 float16 混合，混合精度训练是在尽可能减少精度损失的情况下利用半精度浮点数加速训练。它使用FP16即半精度浮点数存储权重和梯度。在减少占用内存的同时起到了加速训练的效果。参考如下：

https://zhuanlan.zhihu.com/p/79887894

同步BN参考如下：

https://zhuanlan.zhihu.com/p/40496177

train.py代码修改如下：

首先，安装APEX相关库文件，并在代码中引用。

from apex import amp
from apex.parallel import convert_syncbn_model
from apex.parallel import DistributedDataParallel

将模型使用APEX DDP加速：

if args.cuda and args.ddp :
            # self.model = self.model.cuda()
            self.model = self.model.to(device)
            # 同步BN
            self.model = convert_syncbn_model(model)
            # 混合精度
            self.model, self.optimizer = amp.initialize(self.model, self.optimizer, opt_level='O0')
            # 分布数据并行
            self.model = DistributedDataParallel(self.model, delay_allreduce=True)

其中 opt_level 为精度的优化设置，O0（第一个字母是大写字母O）：

O0：纯FP32训练，可以作为accuracy的baseline；
O1：混合精度训练（推荐使用），根据黑白名单自动决定使用FP16（GEMM, 卷积）还是FP32（Softmax）进行计算。
O2：“几乎FP16”混合精度训练，不存在黑白名单，除了Batch norm，几乎都是用FP16计算。
O3：纯FP16训练，很不稳定，但是可以作为speed的baseline；

修改训练迭代时的反向传播实现：

 for i, (image, target) in enumerate(tbar) :
            if self.args.cuda:
                image, target = image.cuda(), target.cuda()
            self.scheduler(self.optimizer, i, epoch, self.best_pred)
            self.optimizer.zero_grad()
            output = self.model(image)
            loss   = self.criterion(output, target)
            #loss   = F.cross_entropy(output, target)
            with amp.scale_loss(loss, self.optimizer) as scaled_loss :
                scaled_loss.backward()

由于分别在两块GPU上做相同的事，所以结果保存和一些打印输出只需要在一块GPU上运行就好：

 if self.args.local_rank == 0 :
     self.saver.save_checkpoint({
         'epoch'      : epoch + 1,
         'state_dict' : self.model.module.state_dict(),
         'optimizer'  : self.optimizer.state_dict(),
         'amp'        : amp.state_dict(),
         'best_pred'  : self.best_pred,
         }, is_best)

增加一些相关配置参数：

# APEX DDP mode
parser.add_argument('--local_rank', default=0, type=int, help='node rank for distributed training')
parser.add_argument('--ddp', type=str, default=True, help='use APEX DDP mode')

# cuda, seed and logging
parser.add_argument('--sync_bn', type=bool, default=True, help='whether to use sync bn (default: auto)')
parser.add_argument('--dsconv', type=bool, default=True, help='whether to use deep separable conv ')

densenet.py文件修改：

增加同步BN选择：

        if sync_bn == True :
            BatchNorm = SynchronizedBatchNorm2d
        else:
            BatchNorm = nn.BatchNorm2d

深度可分离卷积及代码修改

深度可分离卷积相比于标准卷积，在减少计算量的同时并不会大幅降低性能，可参考如下：

https://zhuanlan.zhihu.com/p/92134485

DeepSeperableConv.py增加深度可分离卷积代码实现：

class DeepSeparableConv2d(nn.Module):
    def __init__(self, inplanes, planes, kernel_size=3, stride=1, dilation=2, bias=False, BatchNorm=None):
        super(DeepSeparableConv2d, self).__init__()

        self.conv1 = nn.Conv2d(inplanes, inplanes, kernel_size, stride, 0, dilation,
                               groups=inplanes, bias=bias)
        self.bn = BatchNorm(inplanes)
        self.pointwise = nn.Conv2d(inplanes, planes, 1, 1, 0, 1, 1, bias=bias)

    def forward(self, x):
        x = fixed_padding(x, self.conv1.kernel_size[0], dilation=self.conv1.dilation[0])
        x = self.conv1(x)
        x = self.bn(x)
        x = self.pointwise(x)
        return x

densenet.py添加选择增加深度可分离卷积部分：

        if dsconv == True :
            self.conv2 = DeepSeparableConv2d(nChannels, growthRate, kernel_size=3)
        else :
            self.conv2 = nn.Conv2d(nChannels, growthRate, kernel_size=3, padding=1, bias=False)

保留1*1标准卷积；

detect.py文件修改相应部分。

训练命令如下：

python -m torch.distributed.launch --nproc_per_node=2 train.py  --batch xxx  --lr xxx

--nproc_per_node：每台机器的进程数(显卡)，作者的是单机双显卡，所以是2

凯明权值初始化

He权重初始化详细介绍参考如下链接：

深度学习之参数初始化（二）——Kaiming初始化_凯明初始化-CSDN博客

在densenet.py代码中修改如下：

        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, DeepSeparableConv2d) :
                init.kaiming_normal_(m.weight)
                if m.bias is not None:
                    m.bias.data.fill_(0)
            elif isinstance(m, BatchNorm):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                m.bias.data.zero_()

SWISH和PReLU激活函数

增加了SWISH和PReLU激活函数，详细介绍参考如下：

深入理解深度学习中的激活函数_激活函数的形象解释-CSDN博客

SWISH激活函数代码如下：

class Swish(nn.Module):
    def __init__(self):
        super(Swish, self).__init__()

    def forward(self, x):
        x = x * F.sigmoid(x)
        return x

PReLU和ReLU使用Torch自带的，使用avtive参数选择激活函数，代码如下：

        if active == 'SWISH' :
            act_op = Swish()
        elif active == 'PReLU' :
            act_op = nn.PReLU()
        else :
            act_op = F.relu

测试结果

这个版本主要使用APEX DDP并行加速模式，不支持DP，本来想做成可配置的，发现有点麻烦，所以--ddp参数默认为True。

使用APEX DDP并行加速模式，batch可增大50%，从原来DP的4提升为3*2，如果使用深度可分离卷积，则batch依旧为2*2，训练速度没有仔细计算，目测快了一些；

据说公司接下来会换新丹炉Nvidia RTX 3090，期待一下；

《元龙》作者任怨小说《神工》，再推一波，看得太爽了；

千穹凌帝

关注

15
点赞
踩
9

收藏

觉得还不错? 一键收藏
打赏
0
评论
Pytorch DenseNet训练自己的数据集(续)

Apex 对混合精度训练的过程进行了封装，改两三行配置就可以进行混合精度的训练，从而大幅度降低显存占用，节约运算时间。之前写了一篇，使用pytorch实现DenseNet，完成完整的代码框架，从建立数据集、设置参数、训练网络到推理测试。因此，在上一版的基础上，做些优化改进，设法提升训练速度，降低显存消耗。这个版本主要使用APEX DDP并行加速模式，不支持DP，本来想做成可配置的，发现有点麻烦，所以--ddp参数默认为True。APEX中包含了混合精度，混合精度训练，指代的是。
复制链接

扫一扫