Spark 和 NVIDIA GPU 加速深度学习

随着人们对深度学习（ deep learning ， DL ）兴趣的日益浓厚，越来越多的用户在生产环境中使用 DL 。由于 DL 需要强大的计算能力，开发人员正在利用 gpu 来完成他们的训练和推理工作。

最近，为了更好地统一 Spark 上的 DL 和数据处理，作为 Apache Spark 的一项重大举措的一部分， GPU 成为 Apache Spark 3 . 0 中的一种可调度资源。 Spark 将这些资源请求传递给底层集群管理器。因为这允许您大规模地运行分布式推理，所以它可以帮助加速大数据管道以利用 DL 应用程序。

在 Apache Spark 3 . 0 之前，使用 gpu 是很困难的。您必须手动将 GPU 设备分配给 Spark 作业，并对每个执行器或任务的所有配置进行硬编码，以便在一台机器上使用不同的 GPU 。因为 apachehadoop3 . 1yarn 集群管理器允许不同机器之间的 GPU 协调， Apache Spark 现在可以与它一起工作，帮助将设备安排传递给不同的任务。提交带有 GPU 资源配置发现脚本的应用程序后， Spark 将处理任务之间如何共享 GPU 的分配和协调。

在本教程中，我们将演示如何创建 GPU 机器集群，并使用 Apache Spark 和 Amazon EMR 上的 深层 Java 库（ DJL ） 来利用 Scala 中的大规模图像分类。 DJL 现在提供了一个基于 GPU 的深入学习 Java 包，该包被设计成可以在 Spark 中顺利工作。

如果您对 Scala 和 Java 感兴趣，或者正在寻找将 Java 中的 DL 集成到大数据管道中的解决方案， DJL 提供了一个可行的解决方案。由于 Python 是 DL 最常用的语言，而 Java 是企业开发人员和数据工程师最流行的语言， DJL 的目标是深入学习， Java 开发人员可以使用熟悉的概念和直观的 API 访问的开源工具。 DJL 是建立在现代深度学习框架（ TensorFlow 、 PyTorch 、 Apache MXNet 等）之上的。您可以轻松地使用 DJL 来训练您的模型，或者从各种引擎部署您喜爱的模型，而无需进行任何额外的转换。

设置 Spark 应用程序

有关完整的设置信息，请参阅 gradle 项目设置 。下一节重点介绍一些您应该知道的关键组件。

首先，导入 Spark 依赖项。 Spark SQL 和 ML 库用于存储和处理映像， Spark 依赖项仅在编译时使用，并且由于在运行时提供，因此在打包时被排除在外。当所有东西都打包好后，. jar 任务将它们排除在外。

configurations {
    exclusion
}
dependencies {
    implementation "org.apache.spark:spark-sql_2.12:3.0.1"
    implementation "org.apache.spark:spark-mllib_2.12:3.0.1"
    implementation "org.apache.hadoop:hadoop-hdfs:2.7.4"

    exclusion "org.apache.spark:spark-sql_2.12:3.0.1"
    exclusion "org.apache.spark:spark-mllib_2.12:3.0.1"
    exclusion "org.apache.hadoop:hadoop-hdfs:2.7.4"}
}
jar {
    from {
        (configurations.runtimeClasspath - configurations.exclusion).collect {
            it.isDirectory() ? it : zipTree(it)
        }
    }
}

接下来，导入与 DJL 相关的依赖项。您使用 DJL 和 PyTorch 包。它们提供了 DJL 的核心特性，并加载了一个 DL 引擎来运行以进行推断。此外，您还可以使用

pytorch-native-cu101

在具有 CUDA 10 . 1 的 GPU 上运行。

    implementation platform("ai.djl:bom:0.8.0")
    implementation "ai.djl:api"
    runtimeOnly "ai.djl.pytorch:pytorch-model-zoo"
    runtimeOnly "ai.djl.pytorch:pytorch-native-cu101::linux-x86_64"

加载模型

要在 DJL 中加载模型，请提供承载模型的 URL （ file ://， hdfs ://， s3 ://， https ://）。从该 URL 下载并导入模型。 DJL 还提供了一个强大的动物园模型。 zoo 模型允许您管理预先训练的模型，并在一行中加载它们。内置的动物园模型目前支持 70 多个预先训练和准备使用的模型，这些模型来自 GluonCV 、 HuggingFace 、 TorchHub 和 Keras 。

def loadModel(device : Device): ZooModel[Row, Classifications] = {
    val modelUrl = "https://alpha-djl-demos.s3.amazonaws.com/model/djl-blockrunner/pytorch_resnet18.zip?model_name=traced_resnet18"
    val criteria = Criteria.builder
      .setTypes(classOf[Row], classOf[Classifications])
      .optModelUrls(modelUrl)
      .optTranslator(new MyTranslator())
      .optProgress(new ProgressBar)
      .optDevice(device)
      .build()
     ModelZoo.loadModel(criteria)
}

这里的输入类型是 Spark SQL 中的

Row

。输出类型是分类结果。 MyTranslator 函数执行预处理和后处理工作。加载的模型是来自 torchvision 的预训练 PyTorch ResNet18 模型。

主要逻辑

在下面的代码示例中，

downloadImages

函数下载演示图像并将其存储在 Hadoop 文件系统（

hdfs

）。接下来，

spark.read.format("image")

函数使用 Spark 图像数据源 将图像文件从 HDFS 加载到 Spark

DataFrame

中。在此步骤之后，

mapPartition

获取 GPU 信息。如代码示例所示，

TaskContext.resources()("gpu")

函数存储为此分区分配的 GPU 。这可确保单个设备上的所有 GPU 都得到正确使用。将模型加载到指定的 GPU 后，

predictor.predict(row)

返回 Spark

DataFrame

分区中图像（行）的分类。

def main(args: Array[String]) {

    // download images
    val imagePath = downloadImages(new Path("hdfs:///images"))

    // Spark configuration
    val spark = SparkSession.builder()
      .appName("Image Classification")
      .config(new SparkConf())
      .getOrCreate()
    val df = spark.read.format("image").option("dropInvalid", true).load(imagePath)
    val result = df.select(col("image.*")).mapPartitions(partition => {
      val context = TaskContext.get()
      val gpu = context.resources()("gpu").addresses(0)
      val model = loadModel(Device.gpu(gpu.toInt))
      val predictor = model.newPredictor()
      partition.map(row => {
        predictor.predict(row).toString
      })
    })(Encoders.STRING)
    println(result.collect().mkString("
"))
}

把它包起来

运行

./gradlew jar

将所有内容捆绑到一个 jar 中，并在 Spark 集群中运行。

使用多个 GPU 设置 Spark 群集

由于 Amazon emr6 . 2 . 0 的发布， Spark 3 . 0 在所有 GPU 实例中都可用。

要设置 Spark 群集，请使用 AWS CLI 创建一个包含三个实例的 GPU 群集。要成功运行该命令，必须将

myKey

更改为 EC2 密钥名称。如果预先配置了

--region

选项，也可以将其删除。

aws emr create-cluster \
    --name "Spark cluster" \
    --release-label emr-6.2.0 \
    --region us-east-1 \
    --ebs-root-volume-size 50 \
    --applications Name=Hadoop Name=Spark \
    --ec2-attributes KeyName=myKey \
    --instance-type g3s.xlarge \
    --instance-count 3 \
    --use-default-roles \
    --configurations https://raw.githubusercontent.com/aws-samples/djl-demo/master/aws/emr-distributed-inference/image-classification-gpu/configurations.json 

您可以从 AWS 中提供的各种 GPU 实例中进行选择。此示例使用

g3s.xlarge

实例类型进行测试目的群集设置的总运行时间约为 10 – 15 分钟。

执行 Spark 作业

您可以在 EMR 控制台上或从命令行运行此驻车作业。

下面的命令告诉 Spark 运行一个 Yarn 集群，并设置一个脚本来查找不同设备上的 gpu 。每个任务的 GPU 数量设置为 0 . 5 ，这意味着两个任务共享一个 GPU 。您可能还需要相应地设置 CPU 编号，以确保它们匹配。例如，如果您有一个 8 核 CPU ，并且将

spark.task.cpus

设置为 2 ，这意味着四个任务可以在一台机器上并行运行。要获得最佳性能，请将

spark.task.resource.gpu.amount

设置为 0 . 25 。这允许四个任务共享同一个 GPU 。这有助于最大限度地提高性能，因为 GPU 和 CPU 中的所有核心都已使用。如果没有平衡的设置，一些内核处于空闲状态，这会浪费资源。

spark-submit \
    --master yarn \
    --conf spark.executor.resource.gpu.discoveryScript=/usr/lib/spark/scripts/gpu/getGpusResources.sh \
    --conf spark.worker.resource.gpu.discoveryScript=/usr/lib/spark/scripts/gpu/getGpusResources.sh \
    --conf spark.task.resource.gpu.amount="0.5" \
    --conf spark.task.cpus=2 \
    --conf spark.executor.resource.gpu.amount=1 \
    --conf spark.worker.resource.gpu.amount=1 \
    --class com.examples.ImageClassificationExample \
    build/libs/image-classification-gpu-1.0-SNAPSHOT.jar

这个脚本大约需要 4-6 分钟才能完成，您将得到一个打印输出的推断结果作为输出。