Caffe2 - (十一)ResNet50 Multi-GPU 训练

本文介绍了如何使用Caffe2进行ResNet50模型的多GPU训练。首先概述了ResNet50模型和ImageNet数据集,接着详细讲述了数据集的准备、训练配置。然后,文章逐步解析了创建CNN网络、从DB读取数据、图片变换、创建Residual网络、网络初始化、梯度优化、单GPU及多GPU并行化训练的过程,最后提到了resnet50_trainer.py文件的相关内容。
摘要由CSDN通过智能技术生成

Caffe2 - Multi-GPU 训练

1. 概要

  • ResNet50 model
  • ImageNet 数据集 - 14 million 张图片, 大概需要 300GB SSD 存储空间,2000 个磁盘分片;两张 GPUs 耗时一周.

这里以 ImageNet 中的一部分为例:

  • 640 种 cars 和 640 种 boats 图片集作为训练数据集;
  • 48 种 cars 和 48 种 boats 图片集作为训练数据集;
  • 数据集图片大概 130 MB.

ResNet50 模型训练主要包括:

  • 采用 brew 创建训练网络和测试网络;
  • 采用 model helperCreateDB 来创建图片数据集读取器(database reader);
  • 创建训练函数来基于一张或多张 GPU 进行 ResNet50 模型训练;
  • 创建并行化(parallelized)模型;
  • 循环训练多个 epoches,每个 epoch 中,包括:
    • 对其每个 batch 图片进行模型训练;
    • 运行测试模型;
    • 计算时间,精度,并显示结果.

2. 数据集准备与训练配置

2.1 数据集准备

  • 网络训练前,需要准备训练和测试图片数据集.

    下载 Caffe2 提供的 boats 和 cars 的数据集 resnet_trainer,其选取自 ImageNet,并被转化为 lmdb 格式:

    https://download.caffe2.ai/databases/resnet_trainer.zip
  • 数据加载与python模块导入:

    import numpy as np
    import time
    import os
    
    from caffe2.python import core, workspace, model_helper, net_drawer, memonger, brew
    from caffe2.python import data_parallel_model as dpm
    from caffe2.python.models import resnet
    from caffe2.proto import caffe2_pb2
    workspace.GlobalInit(['caffe2', '--caffe2_log_level=2'])
    
    
    # 训练数据集和测试数据集加载
    
    data_folder = '/path/to/resnet_trainer'
    train_data_db = os.path.join(data_folder, "imagenet_cars_boats_train")
    train_data_db_type = "lmdb"
    
    # 640 cars and 640 boats = 1280
    
    train_data_count = 1280
    test_data_db = os.path.join(data_folder, "imagenet_cars_boats_val")
    test_data_db_type = "lmdb"
    
    # 48 cars and 48 boats = 96
    
    test_data_count = 96
    
    assert os.path.exists(train_data_db)
    assert os.path.exists(test_data_db)

2.2 训练配置

主要是 gpus,batch_size,num_labels,base_learning_rate,stepsize 及 weight_decay 等设置.

# 训练模型用到的 GPUs 数量
# 如, gpus = [0, 1, 2, n] 
gpus = [0]

# Batch size of 32 sums up to roughly 5GB of memory per device
# 每张 GPU 的图片 Batch size 数,每张 GPU 大概需要 5GB 显存
batch_per_device = 32
# 总 Batch size
total_batch_size = batch_per_device * len(gpus)

# 两种 labels: car 和 boat
num_labels = 2

# 初始学习率,缩放因子 - total_batch_size
base_learning_rate = 0.0004 * total_batch_size

# 每 10 个 epochs 改变一次学习率
stepsize = int(10 * train_data_count / total_batch_size)

# Weight decay (L2 regularization)
weight_decay = 1e-4

3. 模型创建与训练

3.1 创建 CNN 网络

采用 Caffe2 Operators - ModelHelper 创建CNN网络:

# model helpe object 仅需要一个参数,即网络名,可以任意命名,主要是对 workspace 网络的引用
# 如:
catos_model = model_helper.ModelHelper(name="catos")

# 创建网络前,清空 workspace
workspace.ResetWorkspace()

3.2 从 DB 读取数据

reader = catos_model.CreateDB(name, db, db_type)

3.3 图片变换

  • Caffe2 编译时需要有 opencv

在实际场景中,图片可能有不同的尺寸(size),长宽比(aspect ratios) 以及旋转角度(orientations),因此训练时需要尽可能的使图片包含更多的情况.

ImageNet 的平均分辨率是 496×387 496 × 387 .

为了便于训练,需要将图片转化为标准尺寸;最直接的做法是简单 resize 到 256×256 256 × 256 ,可参考 Caffe2 - 图像加载与预处理,有对其缺点的介绍.

因此,为了更精确的结果,需要对图片进行合理的 rescale,crop等处理. 虽然也会存在一定的原始图片信息的丢失.

可以围绕图片进行随机裁剪,以得到原始图片的更多变形,扩增训练数据集,增强模型鲁棒性.

如果一张图片中只存在 car 或 boat 的一半,模型最好仍能检测到. 如:
这里写图片描述

图片中仅有 boat 的一半,模型仍得到 50% 的置信度.

Caffe2 提供了 C++ 的图像变换 operator - ImageInput operator,其 Caffe2 的 Python API 使用:

def add_image_input_ops(model):
    # 使用 ImageInput operator 来处理图片
    data, label = model.ImageInput(reader,
                                   ["data", "label"],
                                   batch_size=batch_per_device,
                                   mean=128., # mean: 去除常见 color 均值
                                   std=128., # std: 随机添加对减均值的影响
                                   scale=256, # scale: 将图片 rescale 到通用 size
                                   crop=224, # crop: 裁剪方形图片,提取图片维度信息
                                   is_test=False, # 测试时,不进行图像变换
                                   mirror=1 # 随机进行图片镜像
                                  )
    # 不进行 BP 梯度数值计算
    data = model.StopGradient(data, data)

3.4 创建 Residual 网络

Caffe2 提供了 resnet 的创建函数:from caffe2.python.models import resnet

ResNet50 模型创建:resnet.create_resnet50()函数

create_resnet50(
    model, 
    data, 
    num_input_channels, 
    num_labels, 
    label=None, 
    is_test=False, 
    no_loss=False, 
    no_bias=0, 
    conv1_kernel=7, 
    conv1_stride=2, 
    final_avg_kernel=7
)

create_resnet50_model_ops对该函数的调用:

def create_resnet50_model_ops(model, loss_scale):
    # 创建 Residual 网络
    [softmax, loss] = resnet.create_resnet50(model,
                                             "data",
                                             num_input_channels=3,
                                             num_labels=num_labels,
                                             label="label", )
    prefix = model.net.Proto().name
    loss = model.Scale(loss, prefix + "_loss", scale=loss_scale)
    model.Accuracy([softmax, "label"], prefix + "_accuracy")
    return [loss]

3.5 网络初始化

Caffe2 model helper 对象提供了内在函数,用于采用 BP 算法进行网络学习:

  • AddWeightDecay
  • Iter
  • net.LearningRate
def add_parameter_update_ops(model):
    model.AddWeightDecay(weight_decay)
    iter = model.Iter("iter")
    lr = model.net.LearningRate([iter],
                                "lr",
                                base_lr=base_learning_rate,
                                policy="step",
                                stepsize=stepsize,
                                gamma=0.1, )
    # Momentum SGD update
    for param in model.GetParams():
        param_grad = model.param_to_grad[param]
        param_momentum = model.param_init_net.ConstantFill([param], 
                                                           param + '_momentum', value=0.0)

        # 更新 param_grad and param_momentum in place
        model.net.MomentumSGDUpdate([param_grad, param_momentum, lr, param],
                                    [param_grad, param_momentum, param],
                                    momentum=0.9,
                                    # Nesterov Momentum works slightly better than standard
                                    nesterov=1, )

3.6 梯度优化

如果不采用内存优化,可以减少 batch size,但这里进行了内存优化.

Caffe2 提供了 memonger函数来进行内存优化,重用计算的梯度.

def optimize_gradient_memory(model, loss):
    model.net._net = memonger.share_grad_blobs(model.net,
                                          
评论 9
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值