JVM深度学习实践概述

10月用一整月写完了一篇论文，因此也没能兑现月更的承诺，实属丢人现眼。本文将总结10月以及最近这几天的心得与体会，并希望能够给更多遭遇同样境况的人一些方向。

序

尽管现在在深度学习方面最火热的编程语言是Python，但作为高贵的Java的程序员，使用全动态的Python语言有害于我的纯粹性。因此在前些日子赶论文的时候，一大部分时间都花在了使用JVM进行深度学习上。从一开始的无数次Python入门失败，到今天能够写成文章总结JVM上进行深度学习的经验，我认为这比写论文更有意思，毕竟Java天下第一，在应用领域我还是最支持JVM。

引言

尽管序和引言在中文颇有重复的嫌疑，但即便如此我还是要将他们分开。在本节我将试图阐述为什么Python做机器学习很火，而Java上的深度学习一直不温不火。这里需要强调的是本文所说的深度学习，是通常需要显卡进行训练的任务，而不是传统的机器学习。在传统机器学习领域，我认为Java的Weka库和Oracle的tribuo都已经做的非常优秀了。

关于Python在机器学习领域的火爆，网络上有那么几种说法，我听到的一种说法是CPython对于C系的库调用起来很方便，因为机器学习库底层都是C/C++写的，比较快，而Java通过JavaCPP或者JNI调用普遍比较慢。但是我觉得这种说法并不稳妥，我认为更实际的原因可能是Python好学，毕竟深度学习的核心在于数学，而搞数学的人在以前能写个Matlab就算不错了，而Python又没有什么严格的工程结构，王八蛋都能学会，故因为便于入门，许多搞理论研究的人使用Python构建他们通往实践的路。另一方面Java则没有这么好入门，各种工程设计的模式以及强类型的束缚，以及JVM的深不可测，使得一些Java程序员写了十几年代码也不敢说自己真正了解Java，因此Java才在深度学习领域不温不火吧。

另外值得一说的是Python本身只是一个编程语言规范，而本文所指的Python实际上是其基于C/C++的实现，即CPython，在以前也曾经有过JVM实现的Python，但是其版本止步于2.7，而后来Oracle的GraaleVM也试图打通Python与JVM的隔阂，但就现阶段来说GraalVM上通过Python调用C的API还仍然在开发中。

Support for more extension modules is high priority for us. We are actively building out our support for the Python C API to make extensions such as NumPy, SciPy, Scikit-learn, Pandas, Tensorflow and the like work fully. This work means that some other extensions might also already work, but we're not actively testing other extensions right now and cannot promise anything. Note that to try other extensions on this implementation, you have to download, build, and install them manually for now.

支持更过拓展模块是我们最先要做的。我们正在努力开发Python的C语言API，使得诸如NumPy、SciPy、Scikit-learn、Pandas、Tensorflow之类的框架完全正常工作。这些工作意味着其他一些扩展也可能已经起作用，但是我们现在不积极测试这些扩展，也无法做出任何保证或承诺。 请注意，要尝试此实现的其他扩展，您现在必须手动下载，构建和安装它们。

来源：Github - GraalVM

翻译过来就是GraalVM的Python支持在做，也许能用，但是不稳定。尽管现阶段没有任何应用的可能（我不希望在代码中引入任何不稳定性），但是我还是很看好GraalVM，除了Python，它还有一个Native Image，能够将Java等其支持的语言编译为静态镜像，十分值得期待。

但话说回来，DDL不等人，我们没法等着GraalVM开发完毕再去使用Java调用Python的代码搞深度学习。因此本文将就现阶段可用的，且好用的、符合Java程序员习惯的一些深度学习框架进行介绍与对比。

最后还需要提一点，Python上面的深度学习库依赖于其底层的C++库，原则上可以通过JNI调用他们去做JVM的深度学习，但是很显然这有不是很优雅，我也不认同他是JVM上的深度学习。但后续介绍的一些库也使用，甚至是直接使用的Python深度学习库的C++库，但那样我认为可以是JVM上的深度学习。至于这二者之间的区别，还请读者自行拿捏。

不适用的选项

本节介绍一些对我现有环境不可用的选项，这些选项或许可以在其他环境中正常工作。我的环境是Windows 10.

MXNet Scala Edition

Apache的MXNet在主要面向Python语言之外，还额外的支持了7种语言，其中JVM语言占了3种，分别是Java、Scala和Clojure。其中Java语言的API只能用于现有模型的部署，不能用于训练，Scala语言的API可以进行训练，但只支持CPU，而Clojure版本的API则是对Scala API的再封装。

我选择使用Scala版本，尽管我对SBT一无所知，但还是在学习之初（在Hello World之前）给官方提了一个Issue。后来发现Scala版本的API不支持Windows，只好作罢。

Tensorflow Scala

接着Scala的热度，自然想到的就是Tensorflow在Scala上的非官方封装。对于封装的Scala API本身并没有什么平台之分，而关键在于Tensorflow的动态库，对于预编译版本，并没有Windows的选项，所以只能从Windows上本地编译Tensorflow。而Windows上编译，反正我是不看好。倘若真的需要在Windows上去编译库，我宁愿再装一个Linux双系统也不要安装微软的Visual Studio。所以这个选项也只好作罢。

Tensorflow Java 0.2

如果Scala的路走不通，那就剩下官方的Java API了。Tensorflow的Java API最近发布了0.2版，该版本支持使用Java进行模型训练。但是目前来说文档很少，只有官方的一页，而且Tensorflow在2.0版本全面拥抱Keras，使得Java版本的Tensorflow几乎只能从最底层的做起。尽管有tensorflow-framework起到类似Keras的作用，但目前来说还十分不完善。或许还得等它发展一些时间吧。于是Tensorflow Java也只能作罢。

适用的选项

DeepLearning4J

DL4J是目前我实践最多的一个库，这个库完全不依赖Python库（但底层还是C++实现）的、存粹的商业深度学习库。是的，它是一个商业库，背后的公司曾经是SkyMind，后来变更为konduit.ai。它具有所有其他可以商业使用的库的全部特点：诸如完善的文档，规范的代码组织结构，积极的社区支持等，总的来说用于开发项目十分顺手，十分符合一个Java程序员的使用习惯。

但另一方面，该库对于自定义拓展并不是很友好。目前该框架中有一个SameDiff，旨在模型中引入你自定义的行为，相对于PyTorch之类的针对普通函数的自动微分，Java实现中最困难的还是类型系统。但是好在这种困难并非只针对于这个框架，回想一下Spring框架中为了迎合MVC架构，一个Hello World的网页都要写5个或者更多的类才能实现，在深度学习上只不过是把论文里的公式实现柔和进了继承与覆盖的过程中罢了。即便如此，我还是很开心能够用这个库完成我的论文，比使用Python快乐多了。

为了自定义一个Attention层，首先你需要声明WorkSpace，告诉框架你这个层都需要哪些变量作为参数，这样框架能够根据环境自动在内存或显存种帮你管理这些变量。之后定义运算过程，这个定义只能使用框架提供的有限数学操作，但是你可以使用循环或分支等语句动态的根据参数内容变化最终组合而成的表达式，之后反向传播的时候，框架会根据已知的操作和代码给出的组合自动计算导数等参数。所以其中最困难的地方就是如何面对抽象的对象，将其想象为实际的数据，并利用有限的操作来实现你所要的功能。在这方面我还不是太行，所以至今尚未能实现CRF层（大约半年前），转而使用了Attention。但是随着一篇论文的实践，之后我或许可以做到了。

对于这个框架，如果需要出于研究目的使用的话，其自带的一些东西很可能几乎不能够覆盖到你的需求，你需要自己实现。而实现的过程又极度依赖于Java的开发模式，虽然使用Kotlin在语法上简化了相当一部分工作，但如果对于Java本身并不熟练，要拓展出自己的函数还是十分困难的。而且就目前来说，如何在这个框架上实现一个Transformer Block我还毫无头绪。因此我大概是不会选择使用这个框架进行研究性工作的。但如果需要开发商业或工业的项目，这些项目的重点不在于模型有多新，而是要求程序有多稳，那么这个框架绝对会是我的首选。

Deep Java Library

DJL是亚马逊开发的深度学习库。由于其稀少的文档，使得我在写论文之初选择了DL4J，直到论文完结后我才开始重新审视这个库。

就这个库本身的设计来说，其底层依赖于四个Python深度学习库之一，可以任选Apache MXNet、PyTorch、TensorFlow和ONNX Runtime之一作为底层，而DJL则提供了一个统一的抽象，我选择PyTorch。在使用时只需要配置好对应的maven依赖，完全不需要自己构建什么，maven依赖中会包含对应底层所需要的C++实现库，因此和DL4J一样无痛使用。当然了，前提是你能找到官方的这些maven包的名称。

是的，文档是目前为止这个库最痛苦的地方。你以为在Setup一节中会告诉你如何在Maven或者Gradle中使用这个库？想都别想，只有一段这样的文字：

(Optional) Import using Gradle/Maven wrappers

You use Gradle and Maven wrappers to build the project, so you don't need to install Gradle or Maven. However, you should have basic knowledge about the Gradle or Maven build system.

嗯？那些依赖的包名呢？？？好在这网页有搜索功能，直接搜索maven，找到了相关页面。

总之使用这个库还是挺重度依赖搜索功能的，无论是搜索文档还是搜索谷歌，总归还是有些困难。但是对于库本身来说，因为其底层直接使用各种Python库的依赖，因此在使用上更倾向于已有深度学习库的模样，当然在抽象的过程中会有一些改变，但整体而言不像DL4J那样另起炉灶。在扩展自定义操作方面，可以直接使用LambdaBlock，虽然一样要和DL4J操作一个抽象的对象，但至少在工作量上轻松很多。另外值得一提的是，这个库将模型表示为Block的组合，用起来会和Keras的Sequential十分相似。

昨天摸索了一下午，终于拼凑出来了该库的MNIST程序。

依赖如下（使用Gradle 6.7）：

implementation "org.jetbrains.kotlin:kotlin-stdlib-jdk8:1.4.10"
implementation group: 'ch.qos.logback', name: 'logback-classic', version: '1.2.3'

implementation platform("ai.djl:bom:0.8.0")
implementation 'ai.djl:api:0.8.0'
implementation "ai.djl:basicdataset:0.8.0"
implementation "ai.djl:model-zoo:0.8.0"
runtimeOnly 'ai.djl.pytorch:pytorch-engine:0.8.0'
//    runtimeOnly 'ai.djl.pytorch:pytorch-native-auto:1.6.0'
runtimeOnly 'ai.djl.pytorch:pytorch-native-cu102:1.6.0:win-x86_64'

训练程序如下（使用Kotlin）：

package info.skyblond.djl.demo

import ai.djl.Model
import ai.djl.basicdataset.Mnist
import ai.djl.ndarray.types.Shape
import ai.djl.nn.Activation
import ai.djl.nn.Blocks
import ai.djl.nn.SequentialBlock
import ai.djl.nn.core.Linear
import ai.djl.training.DefaultTrainingConfig
import ai.djl.training.EasyTrain
import ai.djl.training.evaluator.Accuracy
import ai.djl.training.listener.TrainingListener
import ai.djl.training.loss.Loss
import java.nio.file.Files
import java.nio.file.Paths


fun main() {
    Model.newInstance("MNIST").use { model ->
        val batchSize = 32
        val mnist = Mnist.builder()
            .setSampling(batchSize, true)
            .build()

        val input = 28 * 28
        val output = 10

        val hiddenLayerSizes = intArrayOf(128, 64)

        val sequentialBlock = SequentialBlock()

        sequentialBlock.add(Blocks.batchFlattenBlock(input.toLong()))

        for (hiddenSize in hiddenLayerSizes) {
            sequentialBlock.add(Linear.builder().setUnits(hiddenSize.toLong()).build())
            sequentialBlock.add(Activation::relu)
        }

        sequentialBlock.add(Linear.builder().setUnits(output.toLong()).build())
        
        model.block = sequentialBlock

        val config = DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss())
            .addEvaluator(Accuracy())
            .addTrainingListeners(*TrainingListener.Defaults.logging())

        val trainer = model.newTrainer(config)
        trainer.initialize(Shape(1, 28 * 28))
        val epoch = 2

        for (i in 0 until epoch) {
            for (batch in trainer.iterateDataset(mnist)) {
                // During trainBatch, we update the loss and evaluators with the results for the training batch.
                EasyTrain.trainBatch(trainer, batch)
                // Now, we update the model parameters based on the results of the latest trainBatch
                trainer.step()
                // We must make sure to close the batch to ensure all the memory associated with the batch is cleared quickly.
                // If the memory isn't closed after each batch, you will very quickly run out of memory on your GPU
                batch.close()
            }

            // Call the end epoch event for the training listeners now that we are done
            trainer.notifyListeners { listener: TrainingListener ->
                listener.onEpoch(trainer)
            }
        }

        val modelDir = Paths.get("build/mlp")
        Files.createDirectories(modelDir)

        model.setProperty("Epoch", epoch.toString())

        model.save(modelDir, "mlp")
    }
}

推理（即利用模型做预测）程序（使用Kotlin）：

package info.skyblond.djl.demo

import ai.djl.Model
import ai.djl.basicmodelzoo.basic.Mlp
import ai.djl.modality.Classifications
import ai.djl.modality.cv.Image
import ai.djl.modality.cv.ImageFactory
import ai.djl.modality.cv.util.NDImageUtils
import ai.djl.ndarray.NDArray
import ai.djl.ndarray.NDList
import ai.djl.translate.Batchifier
import ai.djl.translate.Translator
import ai.djl.translate.TranslatorContext
import java.nio.file.Paths
import java.util.stream.Collectors
import java.util.stream.IntStream


fun main() {
    val img = ImageFactory.getInstance().fromUrl("https://djl-ai.s3.amazonaws.com/resources/images/0.png");

    val modelDir = Paths.get("build/mlp")
    val model = Model.newInstance("mlp")
    model.block = Mlp(28 * 28, 10, intArrayOf(128, 64))
    model.load(modelDir)

    val translator: Translator<Image, Classifications> = object : Translator<Image, Classifications> {
        override fun processInput(ctx: TranslatorContext, input: Image): NDList {
            // Convert Image to NDArray
            val array: NDArray = input.toNDArray(ctx.ndManager, Image.Flag.GRAYSCALE)
            return NDList(NDImageUtils.toTensor(array))
        }

        override fun processOutput(ctx: TranslatorContext?, list: NDList): Classifications? {
            // Create a Classifications with the output probabilities
            val probabilities = list.singletonOrThrow().softmax(0)
            val classNames = IntStream.range(0, 10).mapToObj { i: Int -> i.toString() }.collect(Collectors.toList())
            return Classifications(classNames, probabilities)
        }

        // The Batchifier describes how to combine a batch together
        // Stacking, the most common batchifier, takes N [X1, X2, ...] arrays to a single [N, X1, X2, ...] array
        override fun getBatchifier(): Batchifier {
            return Batchifier.STACK
        }
    }

    val predictor = model.newPredictor(translator)
    val classifications = predictor.predict(img)
    println(classifications.toString())
}

之后我会更倾向于使用这个库去做一些研究性的东西，从结构上看它会比DL4J更灵活一些。尽管在使用的过程中会有一些阻碍，但好在并不耽误使用，而且代码本身拥有详细的注释，因此一旦引用了Maven依赖，便可以在IDE中通过阅读注释等方法解决半数以上的问题，余下的一半几乎都可以在搜索引擎中得到解决。

小结

DL4J和DJL都能够很方便的通过Maven或Gradle构建深度学习环境，而不需要依赖Python运行环境。至此我认为这两个库已经能够很好的满足在JVM上进行深度学习的需求了，但是除了深度学习本身，关于数据可视化以及深度学习附属的一些领域还差一些。文中没有提到的是DL4J可以通过官方支持使用网页来查看训练过程中的一些数据，诸如当前迭代的分数、更新参数的比例等，有助于调整超参。并且DL4J对于训练和评估的过程十分完备，除了最普通的准确率（Accuracy）之外，还有精确度（Precision）、召回率和F1的计算，而DJL尚未发现有网页UI，而训练过程中的评估，官方也只给了Accuracy一种。但好在DJL也有规范的接口，可以通过继承Evaluator来实现自己的评估器。

尽管还需要一些时间才能更加成熟，无论如何，能够在JVM进行深度学习都是一件很值得高兴的事情。

---

JVM深度学习实践概述 由 天空 Blond 采用 知识共享 署名 - 非商业性使用 - 相同方式共享 4.0 国际 许可协议进行许可。
本许可协议授权之外的使用权限可以从 https://skyblond.info/about.html 处获得。