基于 Julia 的深度学习入门

搬运自我的知乎： https://zhuanlan.zhihu.com/p/142667683

这段时间计算机视觉领域出现了一些使用 Julia 开源的相关工作，要科学合理地对比这些相关工作，储备新的炼丹技巧，笔者不得不开始熟悉 Julia。笔者从周一拿到 Julia 文档开始，这周的试验都是使用 Julia 完成的。这里，打算先说一说笔者的几个感受，帮助大家判断一下自己是否需要着手入坑这门语言：

实用性：★★★★☆

两三年前研究运筹学的时候用 Julia 做最优化问题，感觉比 Cplex 、Matlab 好用。近两年 Julia 开源的深度学习工作逐渐增多，研究的一般是基本问题，在 toy 数据集上跑试验。近期也出现了一些 CV 领域的项目。

生态：★★★☆☆

深度学习库 Flux 和 GPU 计算库 CuArray 基本稳定下来，周边项目更新迅速，比如常用的预训练模型也都可以在 Julia 社区中找到靠谱的库了（如 MetalHead ）。当然，周边项目的快速迭代也会导致一些库动不动就报错（甚至在安装时都要费一番功夫）。另外比较有特点的是，大部分常用的 Python 库都有 PyCall 封装的跟进，实在不行自己用 PyCall 、JavaCall 、Clang 写个胶水层也能用。

易用性：★★★★★

Julia 的语法真的很简单，混合了 Python 和 Matlab，30 分钟入门后续查漏补缺即可。Julia 内置了大量的科学计算方法（符号），确实比 Python 直观和好写了很多。美中不足的是社区现有的代码和官方最佳实践比较少，笔者正在试图在这方面贡献一些工作。

运行速度：★★☆☆☆

运行速度比 PyThon 稍有提高，但是第一次运行需要编译因此调试时体验稍差于 Python。多线程跑崩过系统，GPU 的分布式框架还不太完善。

一、装机必备

在开始之前推荐一些装机必备。考虑到同学们比较熟悉 Python 因此使用 Python 中的 toolbox 进行类比，懒癌患者可以直接装推荐安装的部分：

Julia 对应 Python
Pkg3 对应 pip3
JuliaPro 对应 Anaconda （推荐直接安装这个）
IJulia + Jupyter 对应 IPython + Jupyter （推荐使用）
VSCode Julia 插件对应 VSCode Python 插件
PkgMirrors + 浙大源 对应清华源（推荐使用）

二、炼丹示例

Julia 的语言 Feature 较多，但都比较通俗。因此笔者比较推荐同学们在使用过程中慢慢熟悉（就算你想先慢慢学一个月再去做实验老板也不同意是吧）。如果你实在想先浏览一下基础语法，笔者总结了一个 Notebook，帮助你在 15 分钟内看完并有一个大概印象。

下面笔者总结了 Julia 版的常用 Pipeline，可以帮助同学们理解如何像用 Python + PyTorch 一样简单地使用 Julia 完成深度学习项目。在做实验的时候同学们可以简单复制粘贴，修修改改先跑上。（逃

1. MLP + MNIST 实现一个最小用例

首先，我们先完成一个最小用例，实现在 GPU 上训练一个多层感知器拟合 MNIST，了解基本操作。由于篇幅限制，完整代码请参考并运行 MLP+MNIST。

Flux 是 Julia 中的深度学习库，其完全由 Julia 实现，结构轻量化，是 Julia 中的 PyTorch 。因此首先导入 Flux 备用模型定义和反向传播（训练）。

# 从 Flux 中引入所需组件
using Flux, Flux.Data.MNIST, Statistics
using Flux: onehotbatch, onecold, crossentropy, throttle, params

尽管 Flux 中目前已经实现了 gpu() 方法，但功能有限。所幸 Flux 在 GPU 上的功能基于 CuArrays 实现，可以使用 CUDAapi, CUDAdrv, CUDAnative 来设置 Flux 使用哪个 GPU，或是只使用 CPU 。

using CUDAapi, CUDAdrv, CUDAnative
gpu_id = 1  ## set < 0 for no cuda, >= 0 for using a specific device (if available)

if has_cuda_gpu() && gpu_id >=0
    device!(gpu_id)
    device = Flux.gpu
    @info "Training on GPU-$(gpu_id)"
else
    device = Flux.cpu
    @info "Training on CPU"
end

另外，Flux 目前仍不支持分布式 GPU 训练，要想实现该功能也需要利用上述库写 scatter 和 gather 手动实现。

与 PyTorch 相同，Flux 定义了一个开箱即用的数据集 MNIST 。这里我们调用 MNIST.images() 和 MNIST.labels() 加载数据集和对应的 label，并使用 Flux 中提供的 onehotbatch 对 label 进行 onehot 编码。

imgs = MNIST.images()
labels = onehotbatch(MNIST.labels(), 0:9)

目前，Flux 没有提供数据集切分的函数，因此我们需要手动进行该过程。具体而言，我们使用 partition 对加载进来的数据集进行切分，将每 1000 张图像分为一个 batch，并使用 |> device （遍历每个元素分别执行上文中定义的 device()）全部图像迁移到 GPU 中。

train = [(cat(float.(imgs[i])..., dims = 4), labels[:,i])
         for i in partition(1:60_000, 1000)] |> device

同样，我们选择数据集中前 1000 张图片作为测试数据集，也迁移到 GPU 中。

test_X = cat(float.(MNIST.images(:test)[1:1000])..., dims = 4) |> device
test_y = onehotbatch(MNIST.labels(:test)[1:1000], 0:9) |> device

Flux 中的模型定义与 PyTorch 相似，Chain 取代了 nn.Sequential，Conv/MaxPool/Dense 等 layer 也已经封装好（封装的 cuDNN ）可以直接调用。如下所示，定义模型、损失函数和评估方法只需要三段代码。

model = Chain(
  Conv((2,2), 1=>16, relu),
  MaxPool((2, 2)),
  Conv((2,2), 16=>8, relu),
  MaxPool((2, 2)),
  x -> reshape(x, :, size(x, 4)),
  Dense(288, 10), softmax
) |> device


loss(x, y) = crossentropy(model(x), y)
accuracy(x, y) = mean(onecold(model(x)) .== onecold(y))

Flux 为使用者提供了 Adam 优化器，相比于 PyTorch 的版本，该 Adam 优化器似乎对学习旅更为敏感。如果遇到不收敛的情况可以尝试降低 LR 。后续打算对其 FLux 和 PyTorch 的优化器。和 PyTorch 相似，我们直接使用 ADAM(LR)，定义优化器，使用 train!() 进行训练。

opt = ADAM(0.01)
evalcb() = @show(accuracy(test_X, test_y))

epochs = 5

for i = 1:epochs
    Flux.train!(loss, Flux.params(model), train, opt)
end

值得注意的是 Flux 中构建的图也为动态图，无需考虑计算图的构建，直接定义所需的计算操作就可以了。

进行推断时也如同 Pytorch，可以直接调用模型。如下，从测试集中选择一张图片放入模型，预测所属类别。

using Colors, FileIO, ImageShow

img = test_X[:, :, 1:1, 7:7]

println("Predicted: ", Flux.onecold(model(img |> device)) .- 1)
save("outputs.jpg", collect(test_X[:, :, 1, 7]))

2. VGG + Cifar 封装常用方法 Finetune 模型

在试验和竞赛中，我们通常要对读入图像进行增广；模型也通常是基于某个 pretrained 的模型 Finetune 的，因此接下来我们看如何对这些内容进行封装。由于篇幅限制，这里只说明重要部分，完整代码请参考并运行 VGG+Cifar10。

目前 Flex 和周边的生态还不太完善，图像增强部分的实现实属有限。这里我们参照 pytorch 实现最基本的图像增广的预处理过程。更为丰富的预处理恐怕只能自己编写或是等待官方更新，当然，这也是重新造轮子的好机会～

function resize_smallest_dimension(im, len)
  reduction_factor = len/minimum(size(im)[1:2])
  new_size = size(im)
  new_size = (
      round(Int, size(im,1)*reduction_factor),
      round(Int, size(im,2)*reduction_factor),
  )
  if reduction_factor < 1.0
    # Images.jl's imresize() needs to first lowpass the image, it won't do it for us
    im = imfilter(im, KernelFactors.gaussian(0.75/reduction_factor), Inner())
  end
  return imresize(im, new_size)
end

# Take the len-by-len square of pixels at the center of image `im`
function center_crop(im, len)
  l2 = div(len,2)
  adjust = len % 2 == 0 ? 1 : 0
  return im[div(end,2)-l2:div(end,2)+l2-adjust,div(end,2)-l2:div(end,2)+l2-adjust]
end

function preprocess(im)
  # Resize such that smallest edge is 256 pixels long
  im = resize_smallest_dimension(im, 256)

  # Center-crop to 224x224
  im = center_crop(im, 224)

  # Convert to channel view and normalize (these coefficients taken
  # from PyTorch's ImageNet normalization code)
  μ = [0.485, 0.456, 0.406]
  # the sigma numbers are suspect: they cause the image to go outside of 0..1
  # 1/0.225 = 4.4 effective scale
  σ = [0.229, 0.224, 0.225]
  #im = (channelview(im) .- μ)./σ
  im = (channelview(im) .- μ)

  # Convert from CHW (Image.jl's channel ordering) to WHCN (Flux.jl's ordering)
  # and enforce Float32, as that seems important to Flux
  # result is (224, 224, 3, 1)
  #return Float32.(permutedims(im, (3, 2, 1))[:,:,:,:].*255)  # why
  return Float32.(permutedims(im, (3, 2, 1))[:,:,:,:])
end

这里将 MNIST 的数据集切分方法进行封装，使用 get_processed_data 和 get_test_data 构建训练集合、验证集合和测试集合。

using Metalhead: trainimgs
using Images, ImageMagick

function get_processed_data(args)
    # Fetching the train and validation data and getting them into proper shape	
    X = trainimgs(CIFAR10)
    imgs = [preprocess(X[i].img) for i in 1:40000]
    #onehot encode labels of batch
   
    labels = onehotbatch([X[i].ground_truth.class for i in 1:40000],1:10)

    train_pop = Int((1-args.splitr_)* 40000)
    train = device.([(cat(imgs[i]..., dims = 4), labels[:,i]) for i in partition(1:train_pop, args.batchsize)])
    valset = collect(train_pop+1:40000)
    valX = cat(imgs[valset]..., dims = 4) |> device
    valY = labels[:, valset] |> device

    val = (valX,valY)
    return train, val
end

function get_test_data()
    # Fetch the test data from Metalhead and get it into proper shape.
    test = valimgs(CIFAR10)

    # CIFAR-10 does not specify a validation set so valimgs fetch the testdata instead of testimgs
    testimgs = [preprocess(test[i].img) for i in 1:1000]
    testY = onehotbatch([test[i].ground_truth.class for i in 1:1000], 1:10) |> device
    testX = cat(testimgs..., dims = 4) |> device

    test = (testX,testY)
    return test
end

Julia 中预训练模型库正蓬勃发展，比较成熟的有 Metalhead （类似于 Torchvision ）等。这里我们使用 Metalhead 中提供的模型结构和预训练参数构建 VGG19，并替换后面的层完成当前任务。值得一提的是，目前 EfficientNet 还没有较为优雅的 Julia 封装，实属一大遗憾。

using Metalhead

vgg = VGG19()
model = Chain(vgg.layers[1:end-6],
              Dense(512, 4096, relu),
              Dropout(0.5),
              Dense(4096, 4096, relu),
              Dropout(0.5),
              Dense(4096, 10)) |> device
Flux.trainmode!(model, true)

为了方便试验和记录，我们参照官方实现封装超参数和训练过程。在训练过程中，我们可以定义一个回调函数打印验证集的损失函数：throttle(() -> @show(loss(val...)), args.throttle)。

using Parameters: @with_kw
@with_kw mutable struct Args
    batchsize::Int = 128
    throttle::Int = 10
    lr::Float64 = 5e-5
    epochs::Int = 10
    splitr_::Float64 = 0.1
end

function train(model; kws...)
    # Initialize the hyperparameters
    args = Args(; kws...)
    
    # Load the train, validation data 
    train, val = get_processed_data(args)

    @info("Constructing Model")
    # Defining the loss and accuracy functions

    loss(x, y) = logitcrossentropy(model(x), y)

    ## Training
    # Defining the callback and the optimizer
    evalcb = throttle(() -> @show(loss(val...)), args.throttle)
    opt = ADAM(args.lr)
    @info("Training....")
    # Starting to train models
    Flux.@epochs args.epochs Flux.train!(loss, params(model), train, opt, cb=evalcb)
end

3. ResNet + ImageNet 大型数据集上的标准训练过程

在学会在中小型数据集上完成试验后，我们往往要将试验迁移到大型数据集上。训练过程也会增加很多读取、存储、日志等内容。由于篇幅限制，这里只说明重要部分，完整代码请参考并运行 ResNet+ImageNet。

不同于 PyTorch，目前 Flux 对 Dataset 和 Dataloader 的支持十分有限。官方目前正着力于添加相关功能，不久后可能有相关实现。这里我们模仿 PyTorch 多线程读取数据集并生成 Dataloader 。

struct ImagenetDataset
    # Data we're initialized with
    dataset_root::String
    batch_size::Int
    data_loader::Function

    # Data we calculate once, at startup
    filenames::Vector{String}
    queue_pool::QueuePool

    function ImagenetDataset(dataset_root::String, num_workers::Int, batch_size::Int,
                             data_loader::Function = imagenet_val_data_loader)
        # Scan dataset_root for files
        filenames = filter(f -> endswith(f, ".JPEG"), recursive_readdir(dataset_root))

        @assert !isempty(filenames) "Empty dataset folder!"
        @assert num_workers >= 1 "Must have nonnegative integer number of workers!"
        @assert batch_size >= 1 "Must have nonnegative integer batch size!"

        # Start our worker pool
        @info("Adding $(num_workers) new data workers...")
        queue_pool = QueuePool(num_workers, data_loader, quote
            # The workers need to be able to load images and preprocess them via Metalhead
            using Flux, Images, Metalhead
            include($(@__FILE__))
        end)

        return new(dataset_root, batch_size, data_loader, filenames, queue_pool)
    end
end

# Serialize the arguments needed to recreate this ImagenetDataset
function freeze_args(id::ImagenetDataset)
    return (id.dataset_root, length(id.queue_pool.workers), id.batch_size, id.data_loader)
end
Base.length(id::ImagenetDataset) = div(length(id.filenames),id.batch_size)

mutable struct ImagenetIteratorState
    batch_idx::Int
    job_offset::Int
    
    function ImagenetIteratorState(id::ImagenetDataset)
        @info("Creating IIS with $(length(id.filenames)) images")

        # Build permutation for this iteration
        permutation = shuffle(1:length(id.filenames))

        # Push first job, save value to get job_offset (we know that all jobs
        # within this iteration will be consequtive, so we only save the offset
        # of the first one, and can use that to determine the job ids of every
        # subsequent job:
        filename = joinpath(id.dataset_root, id.filenames[permutation[1]])
        job_offset = push_job!(id.queue_pool, filename)

        # Next, push every other job
        for pidx in permutation[2:end]
            filename = joinpath(id.dataset_root, id.filenames[pidx])
            push_job!(id.queue_pool, filename)
        end
        return new(
            0,
            job_offset,
        )
    end
end

function Base.iterate(id::ImagenetDataset, state=ImagenetIteratorState(id))
    # If we're at the end of this epoch, give up the ghost
    if state.batch_idx > length(id)
        return nothing
    end

    # Otherwise, wait for the next batch worth of jobs to finish on our queue pool
    next_batch_job_ids = state.job_offset .+ (0:(id.batch_size-1)) .+ id.batch_size*state.batch_idx
    # Next, wait for the currently-being-worked-on batch to be done.
    pairs = fetch_result.(Ref(id.queue_pool), next_batch_job_ids)
    state.batch_idx += 1

    # Collate X's and Y's into big tensors:
    X = cat((p[1] for p in pairs)...; dims=ndims(pairs[1][1]))
    Y = cat((p[2] for p in pairs)...; dims=ndims(pairs[1][2]))

    # Return the fruit of our labor
    return (X, Y), state
end

Julia 使用 BSON 实现模型的持久化和读取，速度令人满意。对模型保存和读取进行封装的相关实现如下：

using BSON
using Tracker
using Statistics, Printf
using Flux.Optimise

function save_model(model, filename)
    model_state = Dict(
        :weights => Tracker.data.(params(model))
    )
    open(filename, "w") do io
        BSON.bson(io, model_state)
    end
end

function load_model!(model, filename)
    weights = BSON.load(filename)[:weights]
    Flux.loadparams!(model, weights)
    return model
end

4. DCGAN+Fashion/GCN+Cora 其他网络结构与数据集

近年来 GAN 和 GCN 方兴未艾，只实用 Julia 完成图像分类任务还远远不够。因此笔者正尽可能复现多种类的网络结构和任务。以 GAN 和 GCN 为例，Julia 已经能很好地完成试验目标了。由于篇幅限制，这里只说明重要部分，完整代码请参考并运行 DCGAN+Fashion 和 GCN+Cora。

与 CNN 相同，使用 Flux 可以轻松实现对 DCGAN 的定义。

function Discriminator()
    return Chain(
            Conv((4, 4), 1 => 64; stride = 2, pad = 1),
            x->leakyrelu.(x, 0.2f0),
            Dropout(0.25),
            Conv((4, 4), 64 => 128; stride = 2, pad = 1),
            x->leakyrelu.(x, 0.2f0),
            Dropout(0.25), 
            x->reshape(x, 7 * 7 * 128, :),
            Dense(7 * 7 * 128, 1))
end

function Generator(latent_dim)
    return Chain(
            Dense(latent_dim, 7 * 7 * 256),
            BatchNorm(7 * 7 * 256, relu),
            x->reshape(x, 7, 7, 256, :),
            ConvTranspose((5, 5), 256 => 128; stride = 1, pad = 2),
            BatchNorm(128, relu),
            ConvTranspose((4, 4), 128 => 64; stride = 2, pad = 1),
            BatchNorm(64, relu),
            ConvTranspose((4, 4), 64 => 1, tanh; stride = 2, pad = 1),
            )
end

遵循动态图的反向更新策略，我们只需要像 PyTorch 一样定义对抗损失和对抗训练过程，也较为简单。

function discriminator_loss(real_output, fake_output)
    real_loss = mean(logitbinarycrossentropy.(real_output, 1f0))
    fake_loss = mean(logitbinarycrossentropy.(fake_output, 0f0))
    return real_loss + fake_loss
end

generator_loss(fake_output) = mean(logitbinarycrossentropy.(fake_output, 1f0))

function train_discriminator!(gen, dscr, x, opt_dscr, args)
    noise = randn!(similar(x, (args.latent_dim, args.batch_size))) 
    fake_input = gen(noise)
    ps = Flux.params(dscr)
    # Taking gradient
    loss, back = Flux.pullback(ps) do
        discriminator_loss(dscr(x), dscr(fake_input))
    end
    grad = back(1f0)
    update!(opt_dscr, ps, grad)
    return loss
end

function train_generator!(gen, dscr, x, opt_gen, args)
    noise = randn!(similar(x, (args.latent_dim, args.batch_size))) 
    ps = Flux.params(gen)
    # Taking gradient
    loss, back = Flux.pullback(ps) do
        generator_loss(dscr(gen(noise)))
    end
    grad = back(1f0)
    update!(opt_gen, ps, grad)
    return loss
end

for ep in 1:args.epochs
    @info "Epoch $ep"
    for x in data
        loss_dscr = train_discriminator!(g_model, d_model, x, opt_dscr, args)
        loss_gen = train_generator!(g_model, d_model, x, opt_gen, args)
    end
    train_steps += 1
end

对于其他较为复杂的 CNN 模型，例如 UNet，用户也可以自定义模块的调用过程（类似于 PyTorch 中的 forward ）：

function UNet()
    conv_block = (block1(1, 32), block2(32, 32*2), block2(32*2, 32*4), block2(32*4, 32*8))
    conv_block2 = (block1(32*16, 32*8), block1(32*8, 32*4), block1(32*4, 32*2), block1(32*2, 32))
    bottle = block2(32*8, 32*16)
    upconv_block = (upconv(32*16, 32*8), upconv(32*8, 32*4), upconv(32*4, 32*2), upconv(32*2, 32))
    conv_ = conv(32, 1)
    UNet(conv_block, conv_block2, bottle, upconv_block, conv_)
end

function (u::UNet)(x)
    enc1 = u.conv_block[1](x)
    enc2 = u.conv_block[2](enc1)
    enc3 = u.conv_block[3](enc2)
    enc4 = u.conv_block[4](enc3)
    
    bn = u.bottle(enc4)
    
    dec4 = u.upconv_block[1](bn)
    dec4 = cat(dims=3, dec4, enc4)
    dec4 = u.conv_block2[1](dec4)
    dec3 = u.upconv_block[2](dec4)
    dec3 = cat(dims=3, dec3, enc3)
    dec3 = u.conv_block2[2](dec3)
    dec2 = u.upconv_block[3](dec3)
    dec2 = cat(dims=3, dec2, enc2)
    dec2 = u.conv_block2[3](dec2)
    dec1 = u.upconv_block[4](dec2)
    dec1 = cat(dims=3, dec1, enc1)
    dec1 = u.conv_block2[4](dec1)
    dec1 = u.conv_(dec1)
end

model = UNet()

在 GNN 模型方面，目前较为流行的 GNN 库是 GeometricFlux，但是由于刚刚开源不久，数据读取方面的支持有限。实现应当是参考了 DGL，较为优雅且易于扩展。笔者目前也正在试图基于 LightGraphs 开发一个 GNN 库，主要着力于图的构建和分布式训练部分。

using GeometricFlux

model = Chain(GCNConv(adj_mat, num_features=>hidden, relu),
              Dropout(0.5),
              GCNConv(adj_mat, hidden=>target_catg),
              softmax) |> gpu

三、后记

上述示例代码和讲解均来源于笔者的开源项目 Julia-Deeplearning，目前已有的最佳实践包括：

Julia 教程
- 基础语法
卷积神经网络
生成对抗网络
- DCGAN+Fashion
图卷积网络
- GCN+Cora

由于笔者近期试验较多，因此只能在试验之余偶尔更新。如果同学们有相关工作欢迎 PR 和提 Issue，衷心希望能够抛砖引玉对大家有所帮助～