Tropical Cyclones Detection

The code is adapted from the CMCC use case’s repository.

To know more on the interTwin tropical cyclones detection use case and its DT, please visit the published deliverables, D4.1, D7.1 and D7.3.

Setup env

# After activating the environment
pip install -r requirements.txt

Dataset

If the automatic download from python does not work, try from the command line from within the virtual environment:

gdown https://drive.google.com/drive/folders/1TnmujO4T-8_j4bCxqNe5HEw9njJIIBQD -O data/tmp_data/trainval --folder

For more info visit the gdown repository.

Training

Launch training:

# # ONLY IF tensorflow>=2.16
# export TF_USE_LEGACY_KERAS=1

source ../../.venv-tf/bin/activate
python train.py -p pipeline.yaml

On JSC, the dataset is pre-downloaded and you can use the following command:

# # ONLY IF tensorflow>=2.16
# export TF_USE_LEGACY_KERAS=1

source ../../envAItf_hdfml/bin/activate
python train.py -p pipeline.yaml --data_path /p/project/intertwin/smalldata/cmcc

# Launch a job with SLURM
sbatch startscript.sh

pipeline.yaml

This YAML file defines the pipeline configuration for the CMCC use case.

# General configuration
epochs: 3
micro_batch_size: 32
dataset_url: https://drive.google.com/drive/folders/1TnmujO4T-8_j4bCxqNe5HEw9njJIIBQD #https://drive.google.com/drive/folders/15DEq33MmtRvIpe2bNCg44lnfvEiHcPaf
dataset_root: tmp_cyclones_data
verbose: auto
global_config: null

# Workflows
training_pipeline:
  class_path: itwinai.pipeline.Pipeline
  init_args:
    steps:
      download-step:
        class_path: dataloader.CyclonesDataGetter
        init_args:
          dataset_url: ${dataset_url}
          dataset_root: ${dataset_root}
          global_config: ${global_config}
          patch_type: NEAREST
          shuffle: False
          split_ratio: [0.75, 0.25]
          augment: True
          epochs: ${epochs}
          target_scale: False
          label_no_cyclone: NONE
          aug_type: ONLY_TCS
          experiment: {
            'DRV_VARS_1': ['fg10', 'msl', 't_500', 't_300'],
            'COO_VARS_1': ['patch_cyclone'],
            'MSK_VAR_1': None
          }
          
      training-step:
        class_path: trainer.CyclonesTrainer
        init_args:
          epochs: ${epochs}
          micro_batch_size: ${micro_batch_size}
          global_config: ${global_config}
          network: VGG_V1
          activation: LINEAR
          regularization_strength: NONE
          learning_rate: 0.0001
          loss: MAE
          verbose: ${verbose}

train.py

"""
Training pipeline. To run this script, use the following commands.

On login node:

>>> python train.py -p pipeline.yaml -d

On compute nodes:

>>>  python train.py -p pipeline.yaml

"""

from typing import Dict
import argparse
import logging
from os.path import join
from os import makedirs
from datetime import datetime

# # the mock-0.3.1 dir contains testcase.py, testutils.py & mock.py
from itwinai.parser import ConfigParser, ArgumentParser

from src.macros import PATCH_SIZE, SHAPE


def dynamic_config(args) -> Dict:
    """Generates a configuration with values computed at runtime.

    Args:
        args (argparse.Namespace): arguments parsed from command line.

    Returns:
        Dict: configuration.
    """
    config = {}

    # Paths, Folders
    FORMATTED_DATETIME = str(datetime.now().strftime("%Y-%m-%d_%H-%M-%S"))
    MODEL_BACKUP_DIR = join(args.root_dir, "models/")
    EXPERIMENTS_DIR = join(args.root_dir, "experiments")
    RUN_DIR = join(EXPERIMENTS_DIR, args.run_name +
                   "_" + FORMATTED_DATETIME)
    SCALER_DIR = join(RUN_DIR, "scalers")
    TENSORBOARD_DIR = join(RUN_DIR, "tensorboard")
    CHECKPOINTS_DIR = join(RUN_DIR, "checkpoints")

    # Files
    LOG_FILE = join(RUN_DIR, "run.log")

    # Create folders
    makedirs(EXPERIMENTS_DIR, exist_ok=True)
    makedirs(RUN_DIR, exist_ok=True)
    makedirs(SCALER_DIR, exist_ok=True)
    makedirs(TENSORBOARD_DIR, exist_ok=True)
    makedirs(CHECKPOINTS_DIR, exist_ok=True)

    config = {
        "root_dir": args.root_dir,
        "experiment_dir": EXPERIMENTS_DIR,
        "run_dir": RUN_DIR,
        "scaler_dir": SCALER_DIR,
        "tensorboard_dir": TENSORBOARD_DIR,
        "checkpoints_dir": CHECKPOINTS_DIR,
        "backup_dir": MODEL_BACKUP_DIR,
        "log_file": LOG_FILE,
        "shape": SHAPE,
        "patch_size": PATCH_SIZE,
        # "epochs": args.epochs,
        # "batch_size": args.batch_size
    }

    # initialize logger
    logging.basicConfig(
        format="[%(asctime)s] %(levelname)s : %(message)s",
        level=logging.DEBUG,
        filename=LOG_FILE,
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    return config


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument(
        "-p", "--pipeline", type=str, required=True,
        help='Configuration file to the pipeline to execute.'
    )
    parser.add_argument("-r", "--root_dir", type=str, default='./data')
    parser.add_argument("--data_path", type=str,
                        default='./data/data_path')
    parser.add_argument("-n", "--run_name", default="noname", type=str)
    parser.add_argument(
        '-d', '--download-only',
        action=argparse.BooleanOptionalAction,
        default=False,
        help=('Whether to download only the dataset and exit execution '
              '(suggested on login nodes of HPC systems)')
    )

    args = parser.parse_args()
    global_config = dynamic_config(args)

    # Create parser for the pipeline
    pipe_parser = ConfigParser(
        config=args.pipeline,
        override_keys={
            "dataset_root": args.data_path,
            "global_config": global_config
        }
    )
    pipeline = pipe_parser.parse_pipeline(
        pipeline_nested_key='training_pipeline'
    )

    if args.download_only:
        print('Downloading datasets and exiting...')
        pipeline = pipeline[:1]

    pipeline.execute()

dataloader.py

from typing import List, Dict, Optional
from os import listdir
from os.path import join, exists

import gdown
import numpy as np

from itwinai.components import DataGetter, monitor_exec

from src.macros import (
    PatchType,
    LabelNoCyclone,
    AugmentationType,
)
from src.tfrecords.functions import read_tfrecord_as_tensor
from src.scaling import save_tf_minmax
from src.tfrecords.dataset import eFlowsTFRecordDataset
from src.transform import (
    coo_left_right,
    coo_up_down,
    coo_rot180,
    msk_left_right,
    msk_up_down,
    msk_rot180,
)


class CyclonesDataGetter(DataGetter):
    def __init__(
        self,
        dataset_url: str,
        patch_type: PatchType,
        shuffle: bool,
        split_ratio: List[float],
        augment: bool,
        epochs: int,
        target_scale: bool,
        label_no_cyclone: LabelNoCyclone,
        aug_type: AugmentationType,
        experiment: Dict,
        global_config: Dict,
        shuffle_buffer: Optional[int] = None,
        dataset_root: str = "tmp_data",
        tfrecords_dir: str = "trainval"
    ):
        super().__init__()
        self.save_parameters(**self.locals2params(locals()))
        self.dataset_url = dataset_url
        self.split_ratio = split_ratio
        self.epochs = epochs
        self.target_scale = target_scale
        self.label_no_cyclone = label_no_cyclone.value
        self.shuffle_buffer = shuffle_buffer
        self.aug_type = aug_type.value
        self.patch_type = patch_type.value
        self.augment = augment
        self.global_config = global_config
        self.shuffle = shuffle
        self.dataset_root = dataset_root
        self.tfrecords_dir = tfrecords_dir
        self.drv_vars, self.coo_vars = (
            experiment["DRV_VARS_1"],
            experiment["COO_VARS_1"],
        )
        self.msk_var = (
            None if experiment["MSK_VAR_1"] == "None"
            else experiment["MSK_VAR_!"]
        )
        self.channels = [len(self.drv_vars), len(self.coo_vars)]

        # Shuffle
        if shuffle:
            np.random.shuffle(self.cyclone_files)
            np.random.shuffle(self.adj_files)
            np.random.shuffle(self.random_files)

        # Patches types
        if self.augment:
            if self.msk_var:
                self.aug_fns = {
                    "left_right": msk_left_right,
                    "up_down": msk_up_down,
                    "rot180": msk_rot180,
                }
            else:
                self.aug_fns = {
                    "left_right": coo_left_right,
                    "up_down": coo_up_down,
                    "rot180": coo_rot180,
                }
        else:
            self.aug_fns = {}

        # Parse global config
        self.setup_config(self.global_config)

    def setup_config(self, config: Dict) -> None:
        self.shape = config["shape"]
        root_dir = config["root_dir"]

        # Download data
        if not exists(self.dataset_root):
            gdown.download_folder(
                url=self.dataset_url, quiet=False,
                # verify=False,
                output=self.dataset_root
            )

        # Scalar fields
        self.root_dir = root_dir
        self.tfrecords_path = join(self.dataset_root,
                                   self.tfrecords_dir)
        self.scaler_file = join(config["scaler_dir"], "minmax.tfrecord")

        # get records filenames
        self.cyclone_files = sorted(
            [
                join(self.tfrecords_path, f)
                for f in listdir(self.tfrecords_path)
                if f.endswith(".tfrecord") and f.startswith(
                    PatchType.CYCLONE.value)
            ]
        )
        if self.patch_type == PatchType.NEAREST.value:
            self.adj_files = sorted(
                [
                    join(self.tfrecords_path, f)
                    for f in listdir(self.tfrecords_path)
                    if f.endswith(".tfrecord") and f.startswith(
                        PatchType.NEAREST.value)
                ]
            )
        elif self.patch_type == PatchType.ALLADJACENT.value:
            self.adj_files = sorted(
                [
                    join(self.tfrecords_path, f)
                    for f in listdir(self.tfrecords_path)
                    if f.endswith(".tfrecord")
                    and f.startswith(PatchType.ALLADJACENT.value)
                ]
            )
        self.random_files = sorted(
            [
                join(self.tfrecords_path, f)
                for f in listdir(self.tfrecords_path)
                if f.endswith(".tfrecord") and f.startswith(
                    PatchType.RANDOM.value)
            ]
        )

    def split_files(self, files, ratio):
        n = len(files)
        return (
            files[0: int(ratio[0] * n)],
            files[int(ratio[0] * n): int((ratio[0] + ratio[1]) * n)],
        )

    @monitor_exec
    def execute(self):
        # divide into train, valid and test dataset files
        train_c_fs, valid_c_fs = self.split_files(
            files=self.cyclone_files, ratio=self.split_ratio
        )
        train_a_fs, valid_a_fs = self.split_files(
            files=self.adj_files, ratio=self.split_ratio
        )
        train_r_fs, valid_r_fs = self.split_files(
            files=self.random_files, ratio=self.split_ratio
        )

        # merge all the files together
        train_files = train_c_fs + train_a_fs + train_r_fs
        # valid_files = valid_c_fs + valid_a_fs + valid_r_fs

        # compute scaler on training data
        Xt, _ = read_tfrecord_as_tensor(
            filenames=train_files,
            shape=self.shape,
            drv_vars=self.drv_vars,
            coo_vars=self.coo_vars,
            msk_var=self.msk_var,
        )
        X_scaler = save_tf_minmax(Xt.numpy(), outfile=self.scaler_file)
        scalers = [X_scaler, None]
        Xt = None

        # instantiate training, validation and test sets
        # Contains: (dataset, n_count)
        train_dataset = eFlowsTFRecordDataset(
            cyc_fnames=train_c_fs,
            adj_fnames=train_a_fs,
            rnd_fnames=train_r_fs,
            epochs=self.epochs,
            scalers=scalers,
            target_scale=self.target_scale,
            shape=self.shape,
            label_no_cyclone=self.label_no_cyclone,
            drv_vars=self.drv_vars,
            coo_vars=self.coo_vars,
            msk_var=self.msk_var,
            shuffle_buffer=self.shuffle_buffer,
            aug_fns=self.aug_fns,
            patch_type=self.patch_type,
            aug_type=self.aug_type,
        )
        valid_dataset = eFlowsTFRecordDataset(
            cyc_fnames=valid_c_fs,
            adj_fnames=valid_a_fs,
            rnd_fnames=valid_r_fs,
            epochs=self.epochs,
            scalers=scalers,
            target_scale=self.target_scale,
            shape=self.shape,
            label_no_cyclone=self.label_no_cyclone,
            drv_vars=self.drv_vars,
            coo_vars=self.coo_vars,
            msk_var=self.msk_var,
            shuffle_buffer=self.shuffle_buffer,
            aug_fns=self.aug_fns,
            patch_type=self.patch_type,
            aug_type=self.aug_type,
        )
        return train_dataset, valid_dataset, self.channels

trainer.py

from typing import Dict, Any, Optional, Union
import logging
from os.path import join, exists


import tensorflow as tf
import tensorflow.keras as keras

from itwinai.tensorflow.distributed import get_strategy
from itwinai.tensorflow.trainer import TensorflowTrainer
from itwinai.components import monitor_exec

from src.utils import get_network_config, load_model
from src.callbacks import ProcessBenchmark
from src.macros import (
    Network,
    Losses,
    RegularizationStrength,
    Activation
)


class CyclonesTrainer(TensorflowTrainer):
    strategy: tf.distribute.Strategy
    num_workers: int

    def __init__(
        self,
        network: Network,
        activation: Activation,
        regularization_strength: RegularizationStrength,
        learning_rate: float,
        loss: Losses,
        epochs: int,
        micro_batch_size: int,
        global_config: Dict[str, Any],
        kernel_size: Optional[int] = None,
        model_backup: Optional[str] = None,
        rnd_seed: Optional[int] = None,
        verbose: Union[str, int] = 'auto'
    ):
        super().__init__(
            epochs=epochs,
            micro_batch_size=micro_batch_size,
            rnd_seed=rnd_seed,
            verbose=verbose
        )
        self.save_parameters(**self.locals2params(locals()))
        self.global_config = global_config
        self.model_backup = model_backup
        self.network = network.value
        self.activation = activation.value
        self.kernel_size = kernel_size
        self.regularization_strength, self.regularizer = (
            regularization_strength.value
        )

        # Loss name and learning rate
        self.loss_name, self.loss = loss.value
        self.learning_rate = learning_rate

        # Parse global config
        self.dynamic_config(self.global_config)

    @monitor_exec
    def execute(self, train_data, validation_data, channels) -> None:
        # train_size and valid_size are the number of unique elements
        # in the dataset, before calling tf.data.Dataset.repeat(num_epochs)
        train_dataset, train_size = train_data
        valid_dataset, valid_size = validation_data

        # Batch and distribute datasets among strategy's replica.
        # Each batch is further split among the workers
        dist_train_dataset = self.strategy.experimental_distribute_dataset(
            train_dataset.batch(
                self.macro_batch_size, drop_remainder=True,
                num_parallel_calls=tf.data.AUTOTUNE
            )
        )
        dist_valid_dataset = self.strategy.experimental_distribute_dataset(
            valid_dataset.batch(
                self.macro_batch_size, drop_remainder=True,
                num_parallel_calls=tf.data.AUTOTUNE
            )
        )

        # Inside the strategy load the model, data generators and train
        with self.strategy.scope():
            if not self.model_backup:
                model = get_network_config(
                    network=self.network,
                    patch_size=self.patch_size,
                    activation=self.activation,
                    regularizer=self.regularizer,
                    kernel_size=self.kernel_size,
                    channels=channels,
                )
                logging.debug("New model created")
            else:
                model = load_model(model_fpath=self.best_model_name)
                logging.debug(
                    f"Model loaded from backup at {self.best_model_name}")

            optimizer = keras.optimizers.Adam(learning_rate=self.learning_rate)
            metrics = [keras.metrics.MeanAbsoluteError(name="mae")]
            model.compile(loss=self.loss_name,
                          optimizer=optimizer, metrics=metrics)
        logging.debug("Model compiled")

        # Print model summary to check if model's architecture is correct
        print(model.summary())

        # Compute the steps per epoch for train and valid
        steps_per_epoch = train_size // self.macro_batch_size
        validation_steps = valid_size // self.macro_batch_size

        print("macro_batch_size: ", self.macro_batch_size, flush=True)

        # Train the model
        model.fit(
            dist_train_dataset,
            validation_data=dist_valid_dataset,
            steps_per_epoch=steps_per_epoch,
            validation_steps=validation_steps,
            epochs=self.epochs,
            callbacks=self.callbacks,
        )
        logging.debug("Model trained")

        # Save the best model
        model.save(self.last_model_name)
        logging.debug("Saved training history")

    def dynamic_config(self, config: Dict) -> None:
        """Parse configuration generated at runtime."""
        self.experiment_dir = config["experiment_dir"]
        self.run_dir = config["run_dir"]
        self.patch_size = config["patch_size"]

        # Paths
        CHECKPOINTS_DIR = join(self.run_dir, "checkpoints")

        # files and csvs definition
        CHECKPOINTS_FILEPATH = join(CHECKPOINTS_DIR, "model_{epoch:02d}.keras")
        LOSS_METRICS_HISTORY_CSV = join(
            self.run_dir, "loss_metrics_history.csv")
        BENCHMARK_HISTORY_CSV = join(self.run_dir, "benchmark_history.csv")

        self.callbacks = [
            keras.callbacks.EarlyStopping(
                monitor="val_loss",
                patience=100,
                min_delta=0.0001,
                restore_best_weights=True,
                verbose=1,
                mode="min",
            ),
            keras.callbacks.CSVLogger(LOSS_METRICS_HISTORY_CSV),
            ProcessBenchmark(BENCHMARK_HISTORY_CSV),
            keras.callbacks.ModelCheckpoint(
                filepath=CHECKPOINTS_FILEPATH,
                save_best_only=True,
                monitor="val_loss",
                mode="min",
                save_weights_only=False,
                verbose=1,
            ),
        ]

        # Check if model backup exists
        if self.model_backup is not None and not exists(self.model_backup):
            raise FileNotFoundError("Model backup file not found")
        if self.model_backup:
            self.best_model_name = join(self.model_backup, "best_model.keras")
        self.last_model_name = join(self.run_dir, "last_model.keras")

startscript

#!/bin/bash

# general configuration of the job
#SBATCH --job-name=cyclones
#SBATCH --account=intertwin
#SBATCH --mail-user=
#SBATCH --mail-type=ALL
#SBATCH --output=job.out
#SBATCH --error=job.err
#SBATCH --time=00:30:00

# configure node and process count on the CM
#SBATCH --partition=batch
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=4
#SBATCH --gpus-per-node=4

# SBATCH --exclusive

# gres options have to be disabled for deepv
#SBATCH --gres=gpu:4

set -x
unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY

# load modules
ml --force purge
ml Stages/2024 GCC/12.3.0 OpenMPI CUDA/12 MPI-settings/CUDA Python/3.11 HDF5 PnetCDF libaio mpi4py CMake cuDNN/8.9.5.29-CUDA-12

source ../../envAItf_hdfml/bin/activate

# job info
echo "DEBUG: TIME: $(date)"
echo "DEBUG: EXECUTE: $EXEC"
echo "DEBUG: SLURM_SUBMIT_DIR: $SLURM_SUBMIT_DIR"
echo "DEBUG: SLURM_JOB_ID: $SLURM_JOB_ID"
echo "DEBUG: SLURM_JOB_NODELIST: $SLURM_JOB_NODELIST"
echo "DEBUG: SLURM_NNODES: $SLURM_NNODES"
echo "DEBUG: SLURM_NTASKS: $SLURM_NTASKS"
echo "DEBUG: SLURM_TASKS_PER_NODE: $SLURM_TASKS_PER_NODE"
echo "DEBUG: SLURM_SUBMIT_HOST: $SLURM_SUBMIT_HOST"
echo "DEBUG: SLURMD_NODENAME: $SLURMD_NODENAME"
echo "DEBUG: CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES"
echo "DEBUG: SLURM_NODELIST: $SLURM_NODELIST"
echo

# ONLY IF TENSORFLOW >= 2.16:
# # Using legacy (2.16) version of Keras
# # Latest version with TF (2.16) installs Keras 3.3
# # which returns an error for multi-node execution
# export TF_USE_LEGACY_KERAS=1

# set comm
export CUDA_VISIBLE_DEVICES="0,1,2,3"
export OMP_NUM_THREADS=1
if [ "$SLURM_CPUS_PER_TASK" -gt 0 ] ; then
  export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK
fi

# ON LOGIN NODE download datasets:
# ../../.venv-tf/bin/python train.py -p pipeline.yaml --download-only

# --data_path argument is optional, but on JSC we use the dataset we previously downloaded
srun python train.py -p pipeline.yaml --data_path /p/project/intertwin/smalldata/cmcc

cyclones_vgg.py

import tensorflow as tf


def custom_VGG_V1(patch_size, channels, activation, regularizer):
    model = tf.keras.Sequential()

    model.add(
        tf.keras.layers.Conv2D(
            input_shape=(patch_size, patch_size, channels[0]),
            filters=64,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=64,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=64,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )

    model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

    model.add(
        tf.keras.layers.Conv2D(
            filters=128,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=128,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )

    model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

    model.add(
        tf.keras.layers.Conv2D(
            filters=256,
            kernel_size=(2, 2),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=256,
            kernel_size=(2, 2),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )

    model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

    model.add(
        tf.keras.layers.Conv2D(
            filters=512,
            kernel_size=(2, 2),
            padding="valid",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=512,
            kernel_size=(2, 2),
            padding="valid",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=512,
            kernel_size=(2, 2),
            padding="valid",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )

    model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2)))

    model.add(tf.keras.layers.Flatten())
    model.add(
        tf.keras.layers.Dense(
            units=512, activation=activation, kernel_regularizer=regularizer
        )
    )
    model.add(
        tf.keras.layers.Dense(
            units=256, activation=activation, kernel_regularizer=regularizer
        )
    )
    model.add(
        tf.keras.layers.Dense(
            units=128, activation=activation, kernel_regularizer=regularizer
        )
    )
    model.add(
        tf.keras.layers.Dense(
            units=64, activation=activation, kernel_regularizer=regularizer
        )
    )
    model.add(tf.keras.layers.Dense(channels[1]))

    return model


def custom_VGG_V2(patch_size, channels, activation, regularizer):
    model = tf.keras.Sequential()

    model.add(
        tf.keras.layers.Input(shape=(patch_size, patch_size, channels[0]))
    )

    model.add(
        tf.keras.layers.Conv2D(
            filters=32,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=32,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=32,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )

    model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

    model.add(
        tf.keras.layers.Conv2D(
            filters=64,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=64,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=64,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )

    model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

    model.add(
        tf.keras.layers.Conv2D(
            filters=128,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=128,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=128,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )

    model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

    model.add(
        tf.keras.layers.Conv2D(
            filters=256,
            kernel_size=(2, 2),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=256,
            kernel_size=(2, 2),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=256,
            kernel_size=(2, 2),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )

    model.add(
        tf.keras.layers.Conv2D(
            filters=512,
            kernel_size=(2, 2),
            padding="valid",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=512,
            kernel_size=(2, 2),
            padding="valid",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=512,
            kernel_size=(2, 2),
            padding="valid",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=512,
            kernel_size=(2, 2),
            padding="valid",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )

    model.add(tf.keras.layers.Flatten())

    model.add(
        tf.keras.layers.Dense(
            units=1024, activation=activation, kernel_regularizer=regularizer
        )
    )
    model.add(
        tf.keras.layers.Dense(
            units=512, activation=activation, kernel_regularizer=regularizer
        )
    )
    model.add(
        tf.keras.layers.Dense(
            units=256, activation=activation, kernel_regularizer=regularizer
        )
    )
    model.add(
        tf.keras.layers.Dense(
            units=128, activation=activation, kernel_regularizer=regularizer
        )
    )

    model.add(tf.keras.layers.Dense(channels[1]))

    return model


def custom_VGG_V3(patch_size, channels, activation, regularizer):
    model = tf.keras.Sequential()

    model.add(tf.keras.layers.Input(
        shape=(patch_size, patch_size, channels[0])))

    model.add(
        tf.keras.layers.Conv2D(
            filters=32,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=32,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=32,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )

    model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

    model.add(
        tf.keras.layers.Conv2D(
            filters=64,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=64,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=64,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )

    model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

    model.add(
        tf.keras.layers.Conv2D(
            filters=128,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=128,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=128,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )

    model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

    model.add(
        tf.keras.layers.Conv2D(
            filters=256,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=256,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=256,
            kernel_size=(3, 3),
            padding="same",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )

    model.add(
        tf.keras.layers.Conv2D(
            filters=512,
            kernel_size=(2, 2),
            padding="valid",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=512,
            kernel_size=(2, 2),
            padding="valid",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )

    model.add(
        tf.keras.layers.Conv2D(
            filters=1024,
            kernel_size=(2, 2),
            padding="valid",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )
    model.add(
        tf.keras.layers.Conv2D(
            filters=1024,
            kernel_size=(2, 2),
            padding="valid",
            activation=activation,
            kernel_regularizer=regularizer,
        )
    )

    model.add(tf.keras.layers.Flatten())

    model.add(
        tf.keras.layers.Dense(
            units=1024, activation=activation, kernel_regularizer=regularizer
        )
    )
    model.add(
        tf.keras.layers.Dense(
            units=512, activation=activation, kernel_regularizer=regularizer
        )
    )
    model.add(
        tf.keras.layers.Dense(
            units=512, activation=activation, kernel_regularizer=regularizer
        )
    )
    model.add(
        tf.keras.layers.Dense(
            units=256, activation=activation, kernel_regularizer=regularizer
        )
    )

    model.add(tf.keras.layers.Dense(channels[1]))

    return model


"""
def VGG_V4(patch_size, label_no_cyclone, channels, activation, regularizer):
    model = tf.keras.Sequential()

    model.add(tf.keras.layers.Input(shape=(patch_size, patch_size,
    channels[0])))

    model.add(tf.keras.layers.Conv2D(filters=32, kernel_size=(3,3),
    padding="same", activation=activation, kernel_regularizer=regularizer))
    model.add(tf.keras.layers.Conv2D(filters=32, kernel_size=(3,3),
    padding="same", activation=activation, kernel_regularizer=regularizer))
    model.add(tf.keras.layers.Conv2D(filters=32, kernel_size=(3,3),
    padding="same", activation=activation, kernel_regularizer=regularizer))

    model.add(tf.keras.layers.MaxPooling2D(pool_size=(2,2), strides=(2,2)))

    model.add(tf.keras.layers.Conv2D(filters=64, kernel_size=(3,3),
    padding="same", activation=activation, kernel_regularizer=regularizer))
    model.add(tf.keras.layers.Conv2D(filters=64, kernel_size=(3,3),
    padding="same", activation=activation, kernel_regularizer=regularizer))
    model.add(tf.keras.layers.Conv2D(filters=64, kernel_size=(3,3),
    padding="same", activation=activation, kernel_regularizer=regularizer))

    model.add(tf.keras.layers.MaxPooling2D(pool_size=(2,2), strides=(2,2)))

    model.add(tf.keras.layers.Conv2D(filters=128, kernel_size=(3,3),
    padding="same", activation=activation, kernel_regularizer=regularizer))
    model.add(tf.keras.layers.Conv2D(filters=128, kernel_size=(3,3),
    padding="same", activation=activation, kernel_regularizer=regularizer))
    model.add(tf.keras.layers.Conv2D(filters=128, kernel_size=(3,3),
    padding="same", activation=activation, kernel_regularizer=regularizer))

    model.add(tf.keras.layers.MaxPooling2D(pool_size=(2,2), strides=(2,2)))

    model.add(tf.keras.layers.Conv2D(filters=256, kernel_size=(2,2),
    padding="same", activation=activation, kernel_regularizer=regularizer))
    model.add(tf.keras.layers.Conv2D(filters=256, kernel_size=(2,2),
    padding="same", activation=activation, kernel_regularizer=regularizer))
    model.add(tf.keras.layers.Conv2D(filters=256, kernel_size=(2,2),
    padding="same", activation=activation, kernel_regularizer=regularizer))

    model.add(tf.keras.layers.Conv2D(filters=512, kernel_size=(2,2),
    padding="valid", activation=activation, kernel_regularizer=regularizer))
    model.add(tf.keras.layers.Conv2D(filters=512, kernel_size=(2,2),
    padding="valid", activation=activation, kernel_regularizer=regularizer))
    model.add(tf.keras.layers.Conv2D(filters=512, kernel_size=(2,2),
    padding="valid", activation=activation, kernel_regularizer=regularizer))
    model.add(tf.keras.layers.Conv2D(filters=512, kernel_size=(2,2),
    padding="valid", activation=activation, kernel_regularizer=regularizer))

    model.add(tf.keras.layers.Flatten())

    model.add(tf.keras.layers.Dense(units=1024, activation=activation,
    kernel_regularizer=regularizer))
    model.add(tf.keras.layers.Dense(units=512, activation=activation,
    kernel_regularizer=regularizer))
    model.add(tf.keras.layers.Dense(units=256, activation=activation,
    kernel_regularizer=regularizer))
    model.add(tf.keras.layers.Dense(units=128, activation=activation,
    kernel_regularizer=regularizer))

    model.add(tf.keras.layers.Dense(channels[1]))
    model.add(PositionDiscretizationLayer(label_no_cyclone=label_no_cyclone,
    patch_size=patch_size))

    return model
"""


def ModelV5(patch_size, channels, last_activation, kernel_size=3):
    # kernel initializer
    initializer = tf.random_normal_initializer(0.0, 0.02)

    # input layer
    inputs = tf.keras.layers.Input(shape=(patch_size, patch_size,
channels[0]))

    conv_blocks = [
        ConvBlock(
            filters=32,
            initializer=initializer,
            kernel_size=kernel_size,
            strides=2,
            apply_batchnorm=True,
            apply_dropout=False,
            apply_gaussian_noise=True,
        ),
        ConvBlock(
            filters=64,
            initializer=initializer,
            kernel_size=kernel_size,
            strides=2,
            apply_batchnorm=False,
            apply_dropout=False,
            apply_gaussian_noise=False,
        ),
        ConvBlock(
            filters=128,
            initializer=initializer,
            kernel_size=3,
            strides=2,
            apply_batchnorm=False,
            apply_dropout=True,
            apply_gaussian_noise=False,
        ),
        ConvBlock(
            filters=256,
            initializer=initializer,
            kernel_size=3,
            strides=2,
            apply_batchnorm=False,
            apply_dropout=False,
            apply_gaussian_noise=True,
        ),
        ConvBlock(
            filters=512,
            initializer=initializer,
            kernel_size=3,
            strides=2,
            apply_batchnorm=False,
            apply_dropout=False,
            apply_gaussian_noise=False,
        ),
        ConvBlock(
            filters=1024,
            initializer=initializer,
            kernel_size=3,
            strides=2,
            apply_batchnorm=True,
            apply_dropout=True,
            apply_gaussian_noise=False,
        ),
    ]
    x = inputs
    for block in conv_blocks:
        x = block(x)

    x = tf.keras.layers.Flatten()(x)
    x = tf.keras.layers.Dense(
        units=1024, activation="relu", kernel_initializer=initializer
    )(x)
    x = tf.keras.layers.Dense(
        units=512, activation="relu", kernel_initializer=initializer
    )(x)
    x = tf.keras.layers.Dense(
        units=256, activation="relu", kernel_initializer=initializer
    )(x)
    x = tf.keras.layers.Dense(
        units=128, activation="relu", kernel_initializer=initializer
    )(x)

    outputs = tf.keras.layers.Dense(
        channels[1], activation=last_activation,
kernel_initializer=initializer
    )(x)

    return tf.keras.Model(inputs=inputs, outputs=outputs, name="model_V5")


"""
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense
from tensorflow.keras.models import Sequential
def custom_VGG_V1(patch_size, channels, activation, regularizer):
    model = Sequential()
    #model.add(Conv2D(input_shape=(40,40,len(variables_list)), filters=64,
    kernel_size=(3,3), padding="same", activation=activation,
    kernel_regularizer=regularizer))
    model.add(Conv2D(input_shape=(patch_size,patch_size,channels[0]),
    filters=64, kernel_size=(3,3), padding="same", activation=activation,
    kernel_regularizer=regularizer))
    model.add(Conv2D(filters=64, kernel_size=(3,3), padding="same",
    activation=activation, kernel_regularizer=regularizer))
    model.add(Conv2D(filters=64, kernel_size=(3,3), padding="same",
    activation=activation, kernel_regularizer=regularizer))
    model.add(MaxPooling2D(pool_size=(2,2), strides=(2,2)))
    model.add(Conv2D(filters=128, kernel_size=(3,3), padding="same",
    activation=activation, kernel_regularizer=regularizer))
    model.add(Conv2D(filters=128, kernel_size=(3,3), padding="same",
    activation=activation, kernel_regularizer=regularizer))
    model.add(MaxPooling2D(pool_size=(2,2), strides=(2,2)))
    model.add(Conv2D(filters=256, kernel_size=(2,2), padding="same",
    activation=activation, kernel_regularizer=regularizer))
    model.add(Conv2D(filters=256, kernel_size=(2,2), padding="same",
    activation=activation, kernel_regularizer=regularizer))
    #model.add(Conv2D(filters=256, kernel_size=(2,2), padding="same",
    activation=activation, kernel_regularizer=regularizer))
    model.add(MaxPooling2D(pool_size=(2,2), strides=(2,2)))
    model.add(Conv2D(filters=512, kernel_size=(2,2), padding="valid",
    activation=activation, kernel_regularizer=regularizer))
    model.add(Conv2D(filters=512, kernel_size=(2,2), padding="valid",
    activation=activation, kernel_regularizer=regularizer))
    model.add(Conv2D(filters=512, kernel_size=(2,2), padding="valid",
    activation=activation, kernel_regularizer=regularizer))
    model.add(MaxPooling2D(pool_size=(2,2))) #,strides=(2,2)))
    # model.add(Conv2D(filters=512, kernel_size=(3,3), padding="same",
    activation="relu"))
    # model.add(Conv2D(filters=512, kernel_size=(3,3), padding="same",
    activation="relu"))
    # model.add(Conv2D(filters=512, kernel_size=(3,3), padding="same",
    activation="relu"))
    # model.add(MaxPooling2D(pool_size=(2,2),strides=(2,2)))
    model.add(Flatten())
    model.add(Dense(units=512, activation=activation,
    kernel_regularizer=regularizer))
    model.add(Dense(units=256, activation=activation,
    kernel_regularizer=regularizer))
    model.add(Dense(units=128, activation=activation,
    kernel_regularizer=regularizer))
    model.add(Dense(units=64, activation=activation,
    kernel_regularizer=regularizer))
    model.add(Dense(channels[1]))

    return model
"""