V2EX = way to explore
V2EX 是一个关于分享和探索的地方
现在注册
已注册用户请  登录
tczhangzhi
V2EX  ›  机器学习

基于 Julia 的深度学习入门

  •  
  •   tczhangzhi · 2020-05-23 16:03:09 +08:00 · 2723 次点击
    这是一个创建于 1405 天前的主题,其中的信息可能已经有所发展或是发生改变。

    搬运自我的知乎: 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+FashionGCN+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,目前已有的最佳实践包括:

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

    1 条回复    2020-06-03 23:03:58 +08:00
    formaxin
        1
    formaxin  
       2020-06-03 23:03:58 +08:00 via Android
    从 1 开始有点小难受
    关于   ·   帮助文档   ·   博客   ·   API   ·   FAQ   ·   我们的愿景   ·   实用小工具   ·   1149 人在线   最高记录 6543   ·     Select Language
    创意工作者们的社区
    World is powered by solitude
    VERSION: 3.9.8.5 · 29ms · UTC 22:59 · PVG 06:59 · LAX 15:59 · JFK 18:59
    Developed with CodeLauncher
    ♥ Do have faith in what you're doing.