# Training Deep Neural Network using Data Parallel?

**URL:** <https://discourse.julialang.org/t/training-deep-neural-network-using-data-parallel/75079>\
**Category:** Machine Learning\
**Tags:** parallel, flux\
**Created:** [January 23, 2022, 6:44pm UTC](https://discourse.julialang.org/t/training-deep-neural-network-using-data-parallel/75079 "2022-01-23T18:44:45Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![rkube](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rkube/32/211198_2.png) [@rkube](https://discourse.julialang.org/u/rkube)\
**Post date:** [January 23, 2022, 6:44pm UTC](https://discourse.julialang.org/t/training-deep-neural-network-using-data-parallel/75079/1 "2022-01-23T18:44:45Z")

</div>

Hi,  
I need to train some simpler networks on large datasets and would like to use multiple GPUs for training. Is there a tutorial or writeup on how to do this in Julia? I know about this one in torch: [Multi-GPU Examples — PyTorch Tutorials 1.12.1+cu102 documentation](https://pytorch.org/tutorials/beginner/former_torchies/parallelism_tutorial.html)

Best, Ralph

---

<div class="post-metadata">

**Author:** ![albheim](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/albheim/32/34660_2.png) [@albheim](https://discourse.julialang.org/u/albheim)\
**Post date:** [January 24, 2022, 12:00pm UTC](https://discourse.julialang.org/t/training-deep-neural-network-using-data-parallel/75079/2 "2022-01-24T12:00:20Z")

</div>

Don’t know about anything exactly like the PyTorch example, but you can probably put something together.

This thread could maybe be of interest

> [@Multiple GPUs with Julia](https://discourse.julialang.org/t/multiple-gpus-with-julia/11248):
>
> Has anyone every used multiple GPUs with Flux? If so, did you have to modify your code to run on multiple GPUs?

where the [cuda docs](https://cuda.juliagpu.org/stable/usage/multigpu/) are mentioned, where you can find a little about some basic data handling over gpus.

---

<div class="post-metadata">

**Author:** ![Tomas\_Pevny](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tomas_pevny/32/25466_2.png) [@Tomas\_Pevny](https://discourse.julialang.org/u/Tomas_Pevny)\
**Post date:** [January 24, 2022, 2:09pm UTC](https://discourse.julialang.org/t/training-deep-neural-network-using-data-parallel/75079/3 "2022-01-24T14:09:56Z")

</div>

I have implemented data-parrallel training on CPUs to better utilize multi-threadding in Hierarchical Multiple-instance learning models. It is realtively easy to do by hand. You can check [https://github.com/pevnak/PrayTools.jl/blob/master/src/threadedgrad.jl](https://github.com/pevnak/PrayTools.jl/blob/master/src/threadedgrad.jl) for inspiration. Feel free to contact me for discussion.

Tomas

---

<div class="post-metadata">

**Author:** ![rkube](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rkube/32/211198_2.png) [@rkube](https://discourse.julialang.org/u/rkube)\
**Post date:** [January 24, 2022, 2:55pm UTC](https://discourse.julialang.org/t/training-deep-neural-network-using-data-parallel/75079/4 "2022-01-24T14:55:56Z")

</div>

@Tomas_Pevny Thanks, that’s helpful. I am now sure how to handle Zygote.grads in the distributed setting. Here is my code:

```julia
# MNIST classifier using data-parallel on multiple GPUs

using Distributed
using CUDA

addprocs(length(devices()))

# Include Flux, Zygote, and defines calculation of gradients.
#@everywhere include("/global/homes/r/rkube/source/julia_test/mnist_class_distributed.jl")
@everywhere using Flux
@everywhere using Zygote
@everywhere function grads_now(x, y, mm, pp)
    # xx: Batch of images, size (28)x(28)x(1)x(batch_size)
    # y: One-hot encoded labels of the image, 0-9
    # mm: Flux model
    # pp: Parameters of the flux model
    loss, back = Zygote.pullback(pp) do   
        Flux.Losses.binarycrossentropy(y, mm(x))
    end
    grads = back(1f0)
end

using MLDatasets
using Flux.Data: DataLoader
using Zygote
using Flux

num_gpus = 2 #length(devices())
examples_per_worker = 64;
batch_size = examples_per_worker * num_gpus

# Load data
xtrain, ytrain = MLDatasets.MNIST.traindata(Float32);
xtest, ytest = MLDatasets.MNIST.testdata(Float32);

xtrain = reshape(xtrain, (28, 28, 1, :)) |> gpu;
xtest = reshape(xtest, (28, 28, 1, :)) |> gpu;

# One-hot-encode the labels
ytrain = Flux.onehotbatch(ytrain, 0:9) |> gpu;
ytest = Flux.onehotbatch(ytest, 0:9) |> gpu;

# Batching
train_loader = DataLoader((xtrain, ytrain), batchsize=batch_size, shuffle=true);
test_loader = DataLoader((xtest, ytest), batchsize=batch_size);

# Optimizer
opt = ADAM(0.9);

# Define a simple MLP with sigmoid final layer of size 10 for probability 
model = Chain(Conv((3, 3), 1=>8, relu), 
              Conv((3, 3), 8=>16, relu), 
              Conv((5, 5), 16 =>16, relu), 
              Conv((5, 5), 16=>16, relu), 
              x -> flatten(x), 
              Dense(4096, 10, sigmoid)) |> gpu;

# Extract model params
ps = Flux.params(model);

(x,y) = first(train_loader);

# Split x and y into parts that are sent to each worker
xy_per_worker = [(x[:,:,:, (k-1) * examples_per_worker + 1: k * examples_per_worker],
                  y[:, (k-1) * examples_per_worker + 1: k * examples_per_worker]) for k ∈ 1:num_gpus]

# Calculate gradients over current example distributed on the worker pool
grads_from_workers = pmap(xy -> grads_now(xy[1], xy[2], model, ps), xy_per_worker)
# Average gradients that we just received from the workers
for p in ps
    # Hard-coded for num_gpus=2
    grads_from_workers[1][p] = (grads_from_workers[1][p] + grads_from_workers[2][p]) / 2.0
end

```

The last loop is inspired from your `addgrad` function. Julia throws an error when accessing the elements from `grads_from_workers`:

```julia
julia> for p in ps
           # Hard-coded for num_gpus=2
           grads_from_workers[1][p] = (grads_from_workers[1][p] + grads_from_workers[2][p]) / 2.0
       end
ERROR: KeyError: key [0.26950487 0.09152142 -0.25693536; -0.05512565 -0.0565367 0.26960212; -0.081430614 -0.18807764 -0.07629797;;;; -0.04641332 -0.10099426 0.039984215; 0.19868578 -0.059893914 0.13311733; -0.029146401 -0.15549782 0.13652915;;;; -0.26943153 0.24420884 0.22624107; 0.107807994 0.25059542 0.053223163; -0.19172552 0.11626054 -0.250573;;;; 0.09726549 -0.11461907 -0.21769923; 0.09831371 -0.16952133 -0.2452086; 0.012307492 -0.1203438 -0.20718424;;;; 0.07682779 -0.0055676983 0.008418384; -0.090335764 0.038072642 -0.23139168; -0.1753688 -0.17191374 0.07041458;;;; 0.1565108 0.25166494 0.05122282; 0.057927087 0.14263465 -0.2567391; 0.01940246 -0.09167486 0.2549128;;;; 0.00838662 0.12844574 0.2326228; 0.025148213 0.244805 -0.09040289; -0.05097455 -0.040772133 -0.017126724;;;; -0.13500807 0.00867979 -0.23491067; 0.10669534 0.14557783 -0.16608055; -0.032498617 0.0036792245 0.14004922] not found
Stacktrace:
 [1] getindex
   @ ./iddict.jl:108 [inlined]
 [2] getindex(gs::Zygote.Grads, x::CuArray{Float32, 4, CUDA.Mem.DeviceBuffer})
   @ Zygote ~/.julia/packages/Zygote/kUefI/src/compiler/interface.jl:274
 [3] top-level scope
   @ ./REPL[2]:3

```

When I run this code on a single CPU, accessing the gradients works fine like this:

```julia
julia> gg = grads_now(x, y, model, ps);

julia> gg[ps[1]]
3×3×1×8 Array{Float32, 4}:
[:, :, 1, 1] =
 0.0441429 0.0318292 0.0203954
 0.0292745 0.015554 0.0153562
 0.0194541 0.00176881 -0.00113646

...

```

---

<div class="post-metadata">

**Author:** ![Tomas\_Pevny](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tomas_pevny/32/25466_2.png) [@Tomas\_Pevny](https://discourse.julialang.org/u/Tomas_Pevny)\
**Post date:** [January 24, 2022, 6:24pm UTC](https://discourse.julialang.org/t/training-deep-neural-network-using-data-parallel/75079/5 "2022-01-24T18:24:13Z")

</div>

The best and most effective is to convert them to vector and then restore them back for Grad structure on the master. I have found it very effective, because sending the vector is nice to serialization which is used to send the data between different julia processes.

The cause of error needs a bit deeper understanding of Zygote. Grad contains an `IdDict`, which as a key has a “pointer” to an array and the value is a gradient of that array. As you can imagine, sending pointers over the network does not work well.

You can write a variant of these functions taken from Zygote

```julia
function copy!(gs::Grads, x::AbstractVector, ps)
  i = 0
  for p in ps
    gs[p] .= reshape(x[i+1:i+length(p)], size(p))
    i += length(p)
  end
  gs
end

function copy!(x::AbstractVector, gs::Grads, ps)
  i = 0
  for p in ps
    !haskey(gs, p) && continue
    gs[p] === nothing && continue
    x[i+1:i+length(p)] .= vec(gs[p])
    i += length(p)
  end
  x
end

```

I have adapted them without testing, but you need to pass them an output of `Flux.params(model)` on both sides. Also, do not forget to copy the gradient from gpu if needed using `cpu`.

I hope you got the idea, how to do it. I once had a student, who was implementing some distributed algorithms as a master thesis, but I guess he never released the source code.

Best,

Tomas

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [January 24, 2022, 9:08pm UTC](https://discourse.julialang.org/t/training-deep-neural-network-using-data-parallel/75079/6 "2022-01-24T21:08:19Z")

</div>

Yeah, per my comments in that thread, FluxMPI is probably the most turnkey library available at present. Definitely worth checking out.