# Tutorial on itwinai TorchTrainer adapted for the distributed GAN using MNIST dataset

**Author(s)**: Henry Mutegeki (CERN), Matteo Bunino (CERN), Jarl Sondre Sæther (CERN), Linus Eickhoff (CERN)

The code is adapted from
[this example](https://github.com/pytorch/examples/blob/main/mnist/main.py). Focus is mainly on the `train.py`
file for the distributed GAN use case.

## Setup

First, from the root of this repository, build the environment containing
pytorch and deepspeed.
Refer to the [itwinai installation steps](https://itwinai.readthedocs.io/latest/getting-started/getting_started_with_itwinai.html#install-itwinai).

Then navigate to the project working directory

```bash
cd tutorials/distributed-ml/torch-tutorial-GAN
```

## Distributed training on a single node (interactive)

If you want to use SLURM in interactive mode, do the following:

```bash
# Allocate resources
$ salloc --partition=batch --nodes=1 --account=intertwin  --gres=gpu:4 --time=1:59:00
job ID is XXXX
# Get a shell in the compute node (if using SLURM)
$ srun --jobid XXXX --overlap --pty /bin/bash 
# Now you are inside the compute node

# On JSC, you may need to load some modules
ml --force purge
ml Stages/2024 GCC OpenMPI CUDA/12 MPI-settings/CUDA Python HDF5 PnetCDF libaio mpi4py

# ...before activating the Python environment (adapt this to your env name/path)
source ../../../envAI_hdfml/bin/activate
```

For more info you can refer to
[this](https://itwinai.readthedocs.io/latest/getting-started/slurm.html#interactive-shell-on-a-compute-node)
documentation page.

To launch the training with torch DDP use:

```bash
torchrun --standalone --nnodes=1 --nproc-per-node=gpu train.py --strategy ddp

# Optional -- from a SLURM login node:
srun --jobid XXXX --ntasks-per-node=1 torchrun --standalone --nnodes=1 --nproc-per-node=gpu train.py --strategy ddp
```

To launch the training with Microsoft DeepSpeed use:

```bash
deepspeed train.py -s deepspeed --deepspeed

# Optional -- from a SLURM login node:
srun --jobid XXXX --ntasks-per-node=1 deepspeed train.py --strategy deepspeed
```

## Distributed training with SLURM (batch mode)

You can run your training with SLURM by using the `itwinai` SLURM Builder. Use the
`slurm_config.yaml` file to specify your SLURM parameters and then preview your script
with the following command:

```bash
itwinai generate-slurm -c slurm_config.yaml --no-save-script --no-submit-job
```

If you are happy with the script, you can then run it by omitting `--no-submit-job`:

```bash
itwinai generate-slurm -c slurm_config.yaml --no-save-script
```

If you want to store a copy of the script in a folder, then you can similarly omit
`--no-save-script`:

```bash
itwinai generate-slurm -c slurm_config.yaml
```

## Analyze the logs

Analyze the logs with MLFlow:

```bash
itwinai mlflow-ui --path mllogs/mlflow
```

## Distributed GAN Documentation

This Guide provides a detailed explanation of how a simple Generative Adversarial Network (GAN)
has been adapted to operate within a distributed environment using the `GANTrainer`. This
adaptation enables more efficient training on larger datasets by leveraging distributed
computing resources.

## Overview

A Generative Adversarial Network consists of two key components:

- **Generator (G)**: Generates new data instances.
- **Discriminator (D)**: Evaluates them for authenticity, aiming to distinguish real instances
from the fake ones generated by the Generator.

The training process involves iterative adjustments where the Generator tries to produce data
indistinguishable from actual data, and the Discriminator improves its ability to detect fakes.

## Steps to make a distributed GAN model

The code for all steps can be seen in the attached python file.

### Step 1: Define Model Architecture

Both the Generator and Discriminator are defined using PyTorch's `nn.Module`. The specific
architecture for both includes convolutional layers that are well-suited for processing image
data.

### Step 2: Implement Distributed Training

The `GANTrainer` class extends the custom itwinai `TorchTrainer` class and handles the
initialization of models, optimizers, and the distributed training strategy for the GAN. The
snippet below shows how the GANTrainer is extending the TorchTrainer class and initializing the
parameters. This is essentially done to handle the scenario for the GAN which comprises of two
Neural Network models which is not handled by the TorchTrainer that expects and handles one
model. We also create custom optimizers for the Optimizer and Discriminator GAN models:

### Step 3: Training and Validation Logic

The training alternates between updating the Discriminator using real and generated images and
training the Generator to fool the Discriminator. Validation evaluates the performance of the
Generator in deceiving the Discriminator.

### Step 4: Visualization and Monitoring

Training progress is monitored through visualizations of loss metrics and image samples
generated periodically.

## Takeaways

This readme describes steps taken to adapt a GAN for distributed training, aimed at enhancing
efficiency and scalability for training on large-scale datasets. From this use case we learn
the following:

- itwinai `TorchTrainer` can easily be adapted to different unique use cases like the GAN that
  has two models.
- Training models in a distributed environment may require some level of customization in the
training architecturing but it comes with lots of performance improvements.
- Distributed training for GANs requires a large dataset to reduce the chances of overfitting
to the smaller data splits created during the training phase.
- Always ensure that both the models, data and results for a specific process are accessible on
the same device during training and validation in a distributed environment.
