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

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 helper的 CreateDB 来创建图片数据集读取器(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$.

为了便于训练，需要将图片转化为标准尺寸；最直接的做法是简单 resize 到 $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,