By: Random blog posts about machine learning in Julia

Re-posted from: https://rkube.github.io/julia/gan/2021/11/08/training-gans.html

In order to effectively run machine learning experiments we need a fast
turn-around time for model training. So simply implementing the model is not
the only thing we need to worry about. We also want to be able to change the
hyperparameters in a convenient way. This could either be through a configuration
file or through command line arguments. This post demonstrates how I train
a vanilla GAN on the
MNIST dataset. It is not about GAN theory, for this the original paper by
Goodfellow et al. [[1]] is a good starting point. Instead I focus on how to
structure the code and subtle implementation issues I came across when writing
the code. You can find the current version of the code on github.

Project structure

I am taking a starting point in the vanilla GAN implementation on the
FluxML website. This
implementation works and the trained generator indeed generates images that
look indistinguishable from images belonging to the MNIST dataset.
But how do we arrive there? Why are the learning rates chosen as \(\eta = 2 \times 10^{-4}\)? IS the leakyrelu the optimal activation function or does it perform
on-par with relu in some regime? To answer these questions we need a code that
quickly allows us to change these parameters.

And while we are at it, lets bundle the code together with its dependencies in a
Julia package. This allows us to conveniently a package dependencies to the code.
Taken together, the code and well defined dependencies make the behaviour reproducible. The Julia documentation gives a comprehensive introduction on
packages here.

In order to run the code in the project I first checkout the code from github,
then enter the repository and then execute the runme script:

$ git checkout https://github.com/rkube/mnist_gan.git
$ cd mnist_gan
$ julia --project=. src/runme.jl --activation=ADAM --train_k=8 ...

All packages installed in the project are local to this project and don’t interfere
with packages installed in the general environment. This allows for example to
specify for certain version numbers and will give us producibility of our results.

Code structure

The code is structued as a standard Julia project. The root folder layout looks
like this

├── Manifest.toml
├── Project.toml
├── README.md
└── src
    ├── Manifest.toml
    ├── mnist_gan.jl
    ├── models.jl
    ├── Project.toml
    ├── runme.jl
    └── training.jl

The root folder contains Manifest.toml and Project.toml which include information
about dependencies, versions, package names. More information is given in the
Pkg.jl documentation.

The src folder contains all source codes files. In particular it contains a
mnist_gan.jl file. This is named after the package name and in the simple case
here only twofines the package as a module, includes all other modules and
my two source files

module mnist_gan

using NNlib
using Flux
using Zygote
using Base:Fix2

# All GAN models
include("models.jl")
# Functions used for training
include("training.jl")
end #module

As additional structure I put the models in models.jl and training functions in
training.jl.

Command line arguments

To quickly train the GAN with specific hyperparameters one can either read the
hyperparameters from a configuration file or pass them through the command line.
Here we do the second approach. To comfortably parse command line arguments I’m
using (ArgParse.jl)[https://argparsejl.readthedocs.io/en/latest/argparse.html].

Condensing to only single argument, my code looks like this:

using ArgParse

s = ArgParseSettings()
@add_arg_table s begin
    "--lr_dscr"
        help = "Learning rate for the discriminator. Default = 0.0002"
        arg_type = Float64
        default = 0.0002

args = parse_args(s)

That’s it. Now I can access the single command line arguments via args[lr_dscr].

Logging

Keeping track of the model performance while training is crucial when performing
parameter scans. For the vanilla GAN alone I defined 10 parameters that can be
varied. Letting each parameter assume only two distinct values this allows for
1024 combinations. Julia’s [logging facilities(https://github.com/JuliaLogging)
provide means to systematicallylog model training for a large hyperparameter scan.

In particular, we can use TensorBoardLogger.jl. TensorBoard
provides a visualization of training and includes numerous useful features, such
as visualization of loss curves, displaying of model output images and more. To
use TensorBoardLogger.jl in my code I have to include the module, instantiate
a logger. Then I can easily log my experiments:

# Import the modules
using TensorBoardLogger
...
# Instantiate TensorBoardLogger
# Let's log the hyperparamters of the current run. 
dir_name = join(["$(k)_$(v)" for (k, v) in a])
tb_logger = TBLogger("logs/" * dir_name)
with_logger(tb_logger) do
    @info "hyperparameters" args
end

# Wrap the main training loop in a with clause to enable logging
lossvec_gen = zeros(Float32, args["num_iterations"])
lossvec_dscr = zeros(Float32, ["num_iterations"])

with_logger(tb_logger) do
    for n ∈ args["num_iterations"]e
        # Do machine learning ...
        ...
        # Code to log PNG images to tensorboard, inside the main training loop
        if n % args["output_period"] == 0
            noise = randn(args["latent_dim"], 4) |> gpu;
            fake_img = reshape(generator(noise), 28, 4*28) |> cpu;
            # I need to clip pixel values to [0.0; 1.0]
            fake_img[fake_img .> 1.0] .= 1.0
            fake_img[fake_img .< -1.0] .= -1.0
            fake_img = (fake_img .+ 1.0) .* 0.5
            # 
            log_image(tb_logger, "generatedimage", fake_img, ImageFormat(202))
        end
        # Log the generato and discriminator loss
        @info "test" loss_generator=lossvec_gen[n] loss_discriminator=lossvec_dscr[n]
    end # for
end #  Logger

First, I’m generating a string from all keys and values defined in the command
line argument dictionary. Later this will allow me to filter these arguments.
Then I’m logging the args dictionary, which contains the hyperparameters of
the current experiment. Then I’m generating a fake image using the generator
and log it as well. Here I need to clip the pixel values to [0.0; 1.0]. Since
the Generator is trained on images with pixel values between -1.0 and 1.0 I need
to transform the pixel space. Note that he last argument to the call in
log_image encodes the layout of the fake_img array. I had to look up the
available encodings via

?

Loss functions on-the-fly

To resolve the correct loss function from command line arguments I’m using the
getfield method. To make it a little more convoluted, we also need to distinguish
between loss functions that take an additional, tunable parameterr
like celu, elu, leakyrelu and trelu, and loss functions who do not.
The following code block shows how to map a string that encodes the function name
to the actual function using getfield. To create a closure over an optional
parameter I’m using Fix2. The code below is from models.jl

function get_vanilla_discriminator(args)
    ...
    if args["activation"] in ["celu", "elu", "leakyrelu", "trelu"]
        # Now continue: We want to use Base.Fix2
        act = Fix2(getfield(NNlib, Symbol(args["activation"])), Float32(args["activation_alpha"]))
    else
        act = getfield(NNlib, Symbol(args["activation"]));
    end

    return Chain(Dense(28 * 28, 1024, act), 
        ...);

I found out that can have an impact on performance how I pass the activation
function as an argument to the dense layer. By passing only the function, the
implementation of Dense handles how the activation function is applied to the linear
transformation. This is how it should be. If I manually prescribe how to apply
the broadcast I find slower performance:

julia> d1 = Dense(100, 100, act)
Dense(100, 100, relu)  # 10_100 parameters

julia> @btime d1(randn(Float32, 100, 100));
  163.863 μs (6 allocations: 117.33 KiB)

julia> d2 = Dense(100, 100, x -> act(x))
Dense(100, 100, #5)  # 10_100 parameters

julia> @btime d2(randn(Float32, 100, 100));
  3.041 ms (20016 allocations: 430.30 KiB)

So manually prescribing how to perform the broadcast is about 20 times slower.
Instead, I let the code above return a function that Flux knows how to apply a
broadcast on.

Running a parameter scan

Now we are set up to run a parameter scan. For this I generate runscripts
where I vary my command line arguments. The resulting scripts look like this

#SBATCH things

cd /location/of/the/repo
julia --project=. --lr_dscr=0.0002 --lr_gen=0.0002 --batch_size=8 --num_iterations=10000 --latent_dim=100 --optimizer=ADAM --activation=leakyrelu --activation_alpha=0.2 --train_k=8 --prob_dropout=0.3 --output_period=250

Of course the arguments vary across the scripts. After crunching all the numbers,
the log file directory is populated with the tensorboard log files. The next
blog post will discuss how the results look like and how to pick the best
hyperparameters.

References

[1]
I. Goodfellow et al. Generative Adversarial Networks