分布式AI实践课程简介

分布式AI实践课程目标

本次课程主要介绍使用Intel优化的TensorFlow、Intel MPI Library和Horovod在腾讯云上使用Segmentation Decathlon Dataset中脑肿瘤（BraTS）数据分布式训练一个3D U-Net模型并使用训练出来的模型进行推理。这个模型在仅使用FLAIR（Fluid-attenuated inversion recovery，磁共振成像液体衰减反转恢复序列）Channel的情况下，整个肿瘤的Dice Coefficient能够达到> 0.85。

实验的代码开源在 U-Net Biomedical Image Segmentation with Medical Decathlon Dataset GitHub github.com/shailensobhee/medical-decathlon/tree/master/3D。

本次课程主要使用到的软件环境：

• Ubuntu 18.04, Python 3.6.8
• intel-tensorflow (1.14.0)：github.com/tensorflow/tensorflow
• horovod (0.18.1): github.com/horovod/horovod
• intel-mpi (2019 Update 4): https://software.intel.com/en-us/mpi-library

这些软件已经提前安装在了Intel的分布式AI镜像中，之家直接安装即可使用。

3D U-Net模型

U-Net是卷积神经网络（Convultional Neural Network ）的一种变形，可以利用数据增强（Data Augmentation）可以对一些较少样本的数据进行训练。其网络结构类似于字母U（见右图），因此命名U-Net。

U-Net主要分成两个部分：收缩路径（Contracting Path），扩展路径（Expanding Path）。收缩路径用来捕捉图片的上下文信息（Context Information），扩展路径用来对图片中需要分割出的部分进行定位（Localization）。医学方面相关的数据在获取成本上比其他的数据大非常多，因此U-Net对于医学领域深度学习在较少样本的医学影像的情况有很大的帮助。

医学影像很多时候都是块状的，也就是说是由很多个切片构成一整张图的存在。如果用2D的图像处理模型去处理3D，需要将医学影像图片（包括训练数据和标注数据）切片送进去设计的模型进行训练，这样处理的效率不高。如左下图所示，3D U-Net基于2D U-Net，它和2D U-Net的结构基本是一样的，唯一不同的就是全部2D操作换成了3D。这样子块状图像就不需要单独输入每个切片进行训练，而是可以采取图片整张作为输入到模型中。

Medical Segmentation Decathlon

MSD(Medical Segmentation Decathlon) 提供了10个不同的语义切分的数据和任务可供研究者测试研究机器学习/深度学习算法的普遍可适性。MSD希望实现一个在没有人为干预的情况下，能够解决每个任务的通用算法/模型或者学习系统。使用Medical Segmentation Decathlon的一般步骤：

1. 下载Medical Segmentation Decathlon数据集
2. 实现一个通用的学习算法/模型
3. 对于每一个任务，在没有人为干预的情况下，使用对应的数据集独立训练这个模型（不需要手动去设置每个任务的参数）
4. 用测试数据集测试训练出来的模型，得到切分结果并进行分析

这个实验里，我们只使用Intel优化的TensorFlow和Horovod去分布式训练脑肿瘤的任务模型并推理结果。

腾讯云集群环境搭建

创建服务器实例

• 登录腾讯云账号，在云服务器的实例控制台，点击新建按钮新建服务器实例
• 进入配置的界面。可以快速配置或自定义配置。推荐选择自定义的配置

• 设置主机，可以设置主机的安全组、实例名称和登录方式。设置完成点击立即购买开通实例

• 返回实例控制台，确认实例已经被创建成功

搭建集群环境

• 域名解析。在创建的实例中，选择其中一台作为主节点，其余的则为从节点。在所有的节点中执行vi /etc/hosts，依照下面的格式例出所有实例的内网地址和主机名

• 配置并启动SSH服务。在所有的节点中，创建SSH客户端配置文件执行mkdir /root/.ssh，执行vi /root/.ssh/config修改文件如图所示。

创建SSH远程登录配置文件执行vi /etc/ssh/sshd_config。

创建SSH服务器执行目录执行mkdir -p /var/run/sshd。启动SSH服务器执行/usr/sbin/sshd -D&。

• 配置免密访问。配置从主节点到从节点的免密码访问。在从节点中设置密码执行passwd命令。在主节点中生成SSH秘钥并放到authorized_keys，然后把authorized_keys发送给所有的从节点（需要输入之前设置的从节点密码，具体执行ssh-keygen -t rsa cat ~/.ssh/id_rsa.pub >> ~/.ssh/authorized_keys scp ~/.ssh/authorized_keys root@VM-0-2-ubuntu:~/.ssh/三条命令。检查主节点是否能免密码访问从节点，如能免密登录从节点则配置成功，具体执行ssh root@VM-0-2-ubuntu命令。

运行分布式AI训练和推理

分布式AI训练

在主节点上，进入目标目录执行cd /home/ubuntu/medical-decathlon/3D命令，并使用train_horovod.py在腾讯云多个节点上进行分布式AI训练。

查看train.py里和多节点AI训练通用的训练参数，如下图所示。

默认的参数配置下，原图被分为[Height, Width, Depth] = [144, 144, 144] 的图块且每一个train batch同时处理8个图块，但这需要>=40Gb的内存。如果没有足够大的内存，可设置参数例如--patch_height=16 --patch_width=16 --patch_depth=16 将图块尺寸缩小处理，同时也可以用--bz=4 指定batch size将其缩小。

在进行多节点训练时，需要使用到train_horovod.py 并设置Intel MPI相关的参数和环境变量来达到更好的性能。例如：

mpirun -genv I_MPI_DEBUG=10  -n 3 -hostfile nodelist -ppn 1 -print-rank-map -genv I_MPI_PIN_DOMAIN=socket -genv 
  OMP_NUM_THREADS=8 -genv OMP_PROC_BIND=true -genv KMP_BLOCKTIME=1 python train_horovod.py --bz=4 --
    patch_height=16 --patch_width=16 --patch_depth=16 --epochs=1 --data_path ./Task01_BrainTumour/

• -n：设置有多少MPI进程
• -hostfile：指定所有从节点和主节点的hostname
• -ppn：设置每个节点上有多少进程
• -print-rank-map：报告每个工作进程绑定了哪些节点、socket、core
• I_MPI_PIN_DOMAIN=socket：将worker pin到socket中
• OMP_PROC_BIND=TRUE：将线程绑定到处理器
• KMP_BLOCKTIME=1：设置时间为1ms

更多的Intel MPI参数&环境变量设置详见：https://software.intel.com/en-us/mpi-developer-reference-linux-global-options

训练出来的结果保存在saved_model 目录下，名为3d_unet_decathlon.hdf5在同样的Batch size以及其他条件不变的情况下，使用分布式训练达到了成倍的性能提升。

推理和结果可视化分析

推理命令：

python evaluate_model.py --data_path ./Task01_BrainTumour/ --
saved_model ./saved_model/3d_unet_decathlon.hdf5 

推理的过程中会打印出验证集的平均Dice Coefficient。对于每一个验证集的样本，它会在predictions_directory保存其predictions, MRI images和ground truth masks。

很多应用程序可以用来打开推理出来的Nifti 3D格式的结果，比如Mango (http://ric.uthscsa.edu/mango/mango.html) 。红色部分代表模型预测的脑肿瘤区域，蓝色部分代表实际的脑肿瘤区域。