"""Simplified example of multi-class classification training Model."""

import os
from typing import List, Optional, Dict, Any

from pytorch_lightning.callbacks import EarlyStopping
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.utilities import rank_zero_only
from graphnet.data.dataset.dataset import EnsembleDataset
from graphnet.constants import (
    EXAMPLE_OUTPUT_DIR,
    DATASETS_CONFIG_DIR,
    MODEL_CONFIG_DIR,
)
from graphnet.data.dataloader import DataLoader
from graphnet.data.dataset import Dataset
from graphnet.models import Model
from graphnet.training.callbacks import ProgressBar
from graphnet.utilities.argparse import ArgumentParser
from graphnet.utilities.config import (
    DatasetConfig,
    ModelConfig,
    TrainingConfig,
)
from graphnet.utilities.logging import Logger


def main(
    dataset_config_path: str,
    model_config_path: str,
    gpus: Optional[List[int]],
    max_epochs: int,
    early_stopping_patience: int,
    batch_size: int,
    num_workers: int,
    prediction_names: Optional[List[str]],
    suffix: Optional[str] = None,
    wandb: bool = False,
) -> None:
    """Run example."""
    # Construct Logger
    logger = Logger()

    # Initialise Weights & Biases (W&B) run
    if wandb:
        # Make sure W&B output directory exists
        wandb_dir = "./wandb/"
        os.makedirs(wandb_dir, exist_ok=True)
        wandb_logger = WandbLogger(
            project="example-script",
            entity="graphnet-team",
            save_dir=wandb_dir,
            log_model=True,
        )

    # Build model
    model_config = ModelConfig.load(model_config_path)
    model = Model.from_config(model_config, trust=True)

    # Configuration
    config = TrainingConfig(
        target=[
            target for task in model._tasks for target in task._target_labels
        ],
        early_stopping_patience=early_stopping_patience,
        fit={
            "gpus": gpus,
            "max_epochs": max_epochs,
        },
        dataloader={"batch_size": batch_size, "num_workers": num_workers},
    )

    if suffix is not None:
        archive = os.path.join(EXAMPLE_OUTPUT_DIR, f"train_model_{suffix}")
    else:
        archive = os.path.join(EXAMPLE_OUTPUT_DIR, "train_model")
    run_name = "dynedge_{}_example".format("_".join(config.target))

    # Construct dataloaders
    dataset_config = DatasetConfig.load(dataset_config_path)
    datasets: Dict[str, Any] = Dataset.from_config(
        dataset_config,
    )

    # Construct datasets from multiple selections
    train_dataset = EnsembleDataset(
        [datasets[key] for key in datasets if key.startswith("train")]
    )
    valid_dataset = EnsembleDataset(
        [datasets[key] for key in datasets if key.startswith("valid")]
    )
    test_dataset = EnsembleDataset(
        [datasets[key] for key in datasets if key.startswith("test")]
    )

    # Construct dataloaders
    train_dataloaders = DataLoader(
        train_dataset, shuffle=True, **config.dataloader
    )
    valid_dataloaders = DataLoader(
        valid_dataset, shuffle=False, **config.dataloader
    )
    test_dataloaders = DataLoader(
        test_dataset, shuffle=False, **config.dataloader
    )

    # Log configurations to W&B
    # NB: Only log to W&B on the rank-zero process in case of multi-GPU
    #     training.
    if rank_zero_only == 0:
        wandb_logger.experiment.config.update(config)
        wandb_logger.experiment.config.update(model_config.as_dict())
        wandb_logger.experiment.config.update(dataset_config.as_dict())

    # Build model
    model_config = ModelConfig.load(model_config_path)
    model = Model.from_config(model_config, trust=True)

    # Training model
    callbacks = [
        EarlyStopping(
            monitor="val_loss",
            patience=config.early_stopping_patience,
        ),
        ProgressBar(),
    ]

    model.fit(
        train_dataloaders,
        valid_dataloaders,
        callbacks=callbacks,
        logger=wandb_logger if wandb else None,
        **config.fit,
    )

    # Get predictions
    if isinstance(config.target, str):
        prediction_columns = [
            config.target + "_noise_pred",
            config.target + "_muon_pred",
            config.target + "_neutrino_pred",
        ]
        additional_attributes = [config.target]
    else:
        prediction_columns = [target + "_pred" for target in config.target]
        additional_attributes = config.target

    if prediction_names:
        prediction_columns = prediction_names

    results = model.predict_as_dataframe(
        test_dataloaders,
        prediction_columns=prediction_columns,
        additional_attributes=additional_attributes + ["event_no"],
    )

    # Save predictions and model to file
    db_name = dataset_config.path.split("/")[-1].split(".")[0]
    path = os.path.join(archive, db_name, run_name)
    logger.info(f"Writing results to {path}")
    os.makedirs(path, exist_ok=True)

    results.to_csv(f"{path}/results.csv")
    model.save_state_dict(f"{path}/state_dict.pth")
    model.save(f"{path}/model.pth")


if __name__ == "__main__":
    # Parse command-line arguments
    parser = ArgumentParser(
        description="""
            Train GNN classification model.
            """
    )

    parser.with_standard_arguments(
        (
            "dataset-config",
            os.path.join(
                DATASETS_CONFIG_DIR,
                "training_classification_example_data_sqlite.yml",
            ),
        ),
        (
            "model-config",
            os.path.join(
                MODEL_CONFIG_DIR, "dynedge_PID_classification_example.yml"
            ),
        ),
        "gpus",
        ("max-epochs", 5),
        "early-stopping-patience",
        ("batch-size", 16),
        "num-workers",
    )

    parser.add_argument(
        "--prediction-names",
        nargs="+",
        help="Names of each prediction output feature (default: %(default)s)",
        default=["noise", "muon", "neutrino"],
    )

    parser.add_argument(
        "--suffix",
        type=str,
        help="Name addition to folder (default: %(default)s)",
        default=None,
    )

    args = parser.parse_args()

    main(
        args.dataset_config,
        args.model_config,
        args.gpus,
        args.max_epochs,
        args.early_stopping_patience,
        args.batch_size,
        args.num_workers,
        args.prediction_names,
        args.suffix,
    )

0b3b5f9c31c71f8a966ecd25fea4642ef34ad09e
<mask token>

<mask token>


def main(dataset_config_path: str, model_config_path: str, gpus: Optional[
    List[int]], max_epochs: int, early_stopping_patience: int, batch_size:
    int, num_workers: int, prediction_names: Optional[List[str]], suffix:
    Optional[str]=None, wandb: bool=False) ->None:
    """Run example."""
    logger = Logger()
    if wandb:
        wandb_dir = './wandb/'
        os.makedirs(wandb_dir, exist_ok=True)
        wandb_logger = WandbLogger(project='example-script', entity=
            'graphnet-team', save_dir=wandb_dir, log_model=True)
    model_config = ModelConfig.load(model_config_path)
    model = Model.from_config(model_config, trust=True)
    config = TrainingConfig(target=[target for task in model._tasks for
        target in task._target_labels], early_stopping_patience=
        early_stopping_patience, fit={'gpus': gpus, 'max_epochs':
        max_epochs}, dataloader={'batch_size': batch_size, 'num_workers':
        num_workers})
    if suffix is not None:
        archive = os.path.join(EXAMPLE_OUTPUT_DIR, f'train_model_{suffix}')
    else:
        archive = os.path.join(EXAMPLE_OUTPUT_DIR, 'train_model')
    run_name = 'dynedge_{}_example'.format('_'.join(config.target))
    dataset_config = DatasetConfig.load(dataset_config_path)
    datasets: Dict[str, Any] = Dataset.from_config(dataset_config)
    train_dataset = EnsembleDataset([datasets[key] for key in datasets if
        key.startswith('train')])
    valid_dataset = EnsembleDataset([datasets[key] for key in datasets if
        key.startswith('valid')])
    test_dataset = EnsembleDataset([datasets[key] for key in datasets if
        key.startswith('test')])
    train_dataloaders = DataLoader(train_dataset, shuffle=True, **config.
        dataloader)
    valid_dataloaders = DataLoader(valid_dataset, shuffle=False, **config.
        dataloader)
    test_dataloaders = DataLoader(test_dataset, shuffle=False, **config.
        dataloader)
    if rank_zero_only == 0:
        wandb_logger.experiment.config.update(config)
        wandb_logger.experiment.config.update(model_config.as_dict())
        wandb_logger.experiment.config.update(dataset_config.as_dict())
    model_config = ModelConfig.load(model_config_path)
    model = Model.from_config(model_config, trust=True)
    callbacks = [EarlyStopping(monitor='val_loss', patience=config.
        early_stopping_patience), ProgressBar()]
    model.fit(train_dataloaders, valid_dataloaders, callbacks=callbacks,
        logger=wandb_logger if wandb else None, **config.fit)
    if isinstance(config.target, str):
        prediction_columns = [config.target + '_noise_pred', config.target +
            '_muon_pred', config.target + '_neutrino_pred']
        additional_attributes = [config.target]
    else:
        prediction_columns = [(target + '_pred') for target in config.target]
        additional_attributes = config.target
    if prediction_names:
        prediction_columns = prediction_names
    results = model.predict_as_dataframe(test_dataloaders,
        prediction_columns=prediction_columns, additional_attributes=
        additional_attributes + ['event_no'])
    db_name = dataset_config.path.split('/')[-1].split('.')[0]
    path = os.path.join(archive, db_name, run_name)
    logger.info(f'Writing results to {path}')
    os.makedirs(path, exist_ok=True)
    results.to_csv(f'{path}/results.csv')
    model.save_state_dict(f'{path}/state_dict.pth')
    model.save(f'{path}/model.pth')


<mask token>

<mask token>


def main(dataset_config_path: str, model_config_path: str, gpus: Optional[
    List[int]], max_epochs: int, early_stopping_patience: int, batch_size:
    int, num_workers: int, prediction_names: Optional[List[str]], suffix:
    Optional[str]=None, wandb: bool=False) ->None:
    """Run example."""
    logger = Logger()
    if wandb:
        wandb_dir = './wandb/'
        os.makedirs(wandb_dir, exist_ok=True)
        wandb_logger = WandbLogger(project='example-script', entity=
            'graphnet-team', save_dir=wandb_dir, log_model=True)
    model_config = ModelConfig.load(model_config_path)
    model = Model.from_config(model_config, trust=True)
    config = TrainingConfig(target=[target for task in model._tasks for
        target in task._target_labels], early_stopping_patience=
        early_stopping_patience, fit={'gpus': gpus, 'max_epochs':
        max_epochs}, dataloader={'batch_size': batch_size, 'num_workers':
        num_workers})
    if suffix is not None:
        archive = os.path.join(EXAMPLE_OUTPUT_DIR, f'train_model_{suffix}')
    else:
        archive = os.path.join(EXAMPLE_OUTPUT_DIR, 'train_model')
    run_name = 'dynedge_{}_example'.format('_'.join(config.target))
    dataset_config = DatasetConfig.load(dataset_config_path)
    datasets: Dict[str, Any] = Dataset.from_config(dataset_config)
    train_dataset = EnsembleDataset([datasets[key] for key in datasets if
        key.startswith('train')])
    valid_dataset = EnsembleDataset([datasets[key] for key in datasets if
        key.startswith('valid')])
    test_dataset = EnsembleDataset([datasets[key] for key in datasets if
        key.startswith('test')])
    train_dataloaders = DataLoader(train_dataset, shuffle=True, **config.
        dataloader)
    valid_dataloaders = DataLoader(valid_dataset, shuffle=False, **config.
        dataloader)
    test_dataloaders = DataLoader(test_dataset, shuffle=False, **config.
        dataloader)
    if rank_zero_only == 0:
        wandb_logger.experiment.config.update(config)
        wandb_logger.experiment.config.update(model_config.as_dict())
        wandb_logger.experiment.config.update(dataset_config.as_dict())
    model_config = ModelConfig.load(model_config_path)
    model = Model.from_config(model_config, trust=True)
    callbacks = [EarlyStopping(monitor='val_loss', patience=config.
        early_stopping_patience), ProgressBar()]
    model.fit(train_dataloaders, valid_dataloaders, callbacks=callbacks,
        logger=wandb_logger if wandb else None, **config.fit)
    if isinstance(config.target, str):
        prediction_columns = [config.target + '_noise_pred', config.target +
            '_muon_pred', config.target + '_neutrino_pred']
        additional_attributes = [config.target]
    else:
        prediction_columns = [(target + '_pred') for target in config.target]
        additional_attributes = config.target
    if prediction_names:
        prediction_columns = prediction_names
    results = model.predict_as_dataframe(test_dataloaders,
        prediction_columns=prediction_columns, additional_attributes=
        additional_attributes + ['event_no'])
    db_name = dataset_config.path.split('/')[-1].split('.')[0]
    path = os.path.join(archive, db_name, run_name)
    logger.info(f'Writing results to {path}')
    os.makedirs(path, exist_ok=True)
    results.to_csv(f'{path}/results.csv')
    model.save_state_dict(f'{path}/state_dict.pth')
    model.save(f'{path}/model.pth')


if __name__ == '__main__':
    parser = ArgumentParser(description=
        """
            Train GNN classification model.
            """)
    parser.with_standard_arguments(('dataset-config', os.path.join(
        DATASETS_CONFIG_DIR,
        'training_classification_example_data_sqlite.yml')), (
        'model-config', os.path.join(MODEL_CONFIG_DIR,
        'dynedge_PID_classification_example.yml')), 'gpus', ('max-epochs', 
        5), 'early-stopping-patience', ('batch-size', 16), 'num-workers')
    parser.add_argument('--prediction-names', nargs='+', help=
        'Names of each prediction output feature (default: %(default)s)',
        default=['noise', 'muon', 'neutrino'])
    parser.add_argument('--suffix', type=str, help=
        'Name addition to folder (default: %(default)s)', default=None)
    args = parser.parse_args()
    main(args.dataset_config, args.model_config, args.gpus, args.max_epochs,
        args.early_stopping_patience, args.batch_size, args.num_workers,
        args.prediction_names, args.suffix)

<mask token>
import os
from typing import List, Optional, Dict, Any
from pytorch_lightning.callbacks import EarlyStopping
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.utilities import rank_zero_only
from graphnet.data.dataset.dataset import EnsembleDataset
from graphnet.constants import EXAMPLE_OUTPUT_DIR, DATASETS_CONFIG_DIR, MODEL_CONFIG_DIR
from graphnet.data.dataloader import DataLoader
from graphnet.data.dataset import Dataset
from graphnet.models import Model
from graphnet.training.callbacks import ProgressBar
from graphnet.utilities.argparse import ArgumentParser
from graphnet.utilities.config import DatasetConfig, ModelConfig, TrainingConfig
from graphnet.utilities.logging import Logger


def main(dataset_config_path: str, model_config_path: str, gpus: Optional[
    List[int]], max_epochs: int, early_stopping_patience: int, batch_size:
    int, num_workers: int, prediction_names: Optional[List[str]], suffix:
    Optional[str]=None, wandb: bool=False) ->None:
    """Run example."""
    logger = Logger()
    if wandb:
        wandb_dir = './wandb/'
        os.makedirs(wandb_dir, exist_ok=True)
        wandb_logger = WandbLogger(project='example-script', entity=
            'graphnet-team', save_dir=wandb_dir, log_model=True)
    model_config = ModelConfig.load(model_config_path)
    model = Model.from_config(model_config, trust=True)
    config = TrainingConfig(target=[target for task in model._tasks for
        target in task._target_labels], early_stopping_patience=
        early_stopping_patience, fit={'gpus': gpus, 'max_epochs':
        max_epochs}, dataloader={'batch_size': batch_size, 'num_workers':
        num_workers})
    if suffix is not None:
        archive = os.path.join(EXAMPLE_OUTPUT_DIR, f'train_model_{suffix}')
    else:
        archive = os.path.join(EXAMPLE_OUTPUT_DIR, 'train_model')
    run_name = 'dynedge_{}_example'.format('_'.join(config.target))
    dataset_config = DatasetConfig.load(dataset_config_path)
    datasets: Dict[str, Any] = Dataset.from_config(dataset_config)
    train_dataset = EnsembleDataset([datasets[key] for key in datasets if
        key.startswith('train')])
    valid_dataset = EnsembleDataset([datasets[key] for key in datasets if
        key.startswith('valid')])
    test_dataset = EnsembleDataset([datasets[key] for key in datasets if
        key.startswith('test')])
    train_dataloaders = DataLoader(train_dataset, shuffle=True, **config.
        dataloader)
    valid_dataloaders = DataLoader(valid_dataset, shuffle=False, **config.
        dataloader)
    test_dataloaders = DataLoader(test_dataset, shuffle=False, **config.
        dataloader)
    if rank_zero_only == 0:
        wandb_logger.experiment.config.update(config)
        wandb_logger.experiment.config.update(model_config.as_dict())
        wandb_logger.experiment.config.update(dataset_config.as_dict())
    model_config = ModelConfig.load(model_config_path)
    model = Model.from_config(model_config, trust=True)
    callbacks = [EarlyStopping(monitor='val_loss', patience=config.
        early_stopping_patience), ProgressBar()]
    model.fit(train_dataloaders, valid_dataloaders, callbacks=callbacks,
        logger=wandb_logger if wandb else None, **config.fit)
    if isinstance(config.target, str):
        prediction_columns = [config.target + '_noise_pred', config.target +
            '_muon_pred', config.target + '_neutrino_pred']
        additional_attributes = [config.target]
    else:
        prediction_columns = [(target + '_pred') for target in config.target]
        additional_attributes = config.target
    if prediction_names:
        prediction_columns = prediction_names
    results = model.predict_as_dataframe(test_dataloaders,
        prediction_columns=prediction_columns, additional_attributes=
        additional_attributes + ['event_no'])
    db_name = dataset_config.path.split('/')[-1].split('.')[0]
    path = os.path.join(archive, db_name, run_name)
    logger.info(f'Writing results to {path}')
    os.makedirs(path, exist_ok=True)
    results.to_csv(f'{path}/results.csv')
    model.save_state_dict(f'{path}/state_dict.pth')
    model.save(f'{path}/model.pth')


if __name__ == '__main__':
    parser = ArgumentParser(description=
        """
            Train GNN classification model.
            """)
    parser.with_standard_arguments(('dataset-config', os.path.join(
        DATASETS_CONFIG_DIR,
        'training_classification_example_data_sqlite.yml')), (
        'model-config', os.path.join(MODEL_CONFIG_DIR,
        'dynedge_PID_classification_example.yml')), 'gpus', ('max-epochs', 
        5), 'early-stopping-patience', ('batch-size', 16), 'num-workers')
    parser.add_argument('--prediction-names', nargs='+', help=
        'Names of each prediction output feature (default: %(default)s)',
        default=['noise', 'muon', 'neutrino'])
    parser.add_argument('--suffix', type=str, help=
        'Name addition to folder (default: %(default)s)', default=None)
    args = parser.parse_args()
    main(args.dataset_config, args.model_config, args.gpus, args.max_epochs,
        args.early_stopping_patience, args.batch_size, args.num_workers,
        args.prediction_names, args.suffix)

"""Simplified example of multi-class classification training Model."""

import os
from typing import List, Optional, Dict, Any

from pytorch_lightning.callbacks import EarlyStopping
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.utilities import rank_zero_only
from graphnet.data.dataset.dataset import EnsembleDataset
from graphnet.constants import (
    EXAMPLE_OUTPUT_DIR,
    DATASETS_CONFIG_DIR,
    MODEL_CONFIG_DIR,
)
from graphnet.data.dataloader import DataLoader
from graphnet.data.dataset import Dataset
from graphnet.models import Model
from graphnet.training.callbacks import ProgressBar
from graphnet.utilities.argparse import ArgumentParser
from graphnet.utilities.config import (
    DatasetConfig,
    ModelConfig,
    TrainingConfig,
)
from graphnet.utilities.logging import Logger


def main(
    dataset_config_path: str,
    model_config_path: str,
    gpus: Optional[List[int]],
    max_epochs: int,
    early_stopping_patience: int,
    batch_size: int,
    num_workers: int,
    prediction_names: Optional[List[str]],
    suffix: Optional[str] = None,
    wandb: bool = False,
) -> None:
    """Run example."""
    # Construct Logger
    logger = Logger()

    # Initialise Weights & Biases (W&B) run
    if wandb:
        # Make sure W&B output directory exists
        wandb_dir = "./wandb/"
        os.makedirs(wandb_dir, exist_ok=True)
        wandb_logger = WandbLogger(
            project="example-script",
            entity="graphnet-team",
            save_dir=wandb_dir,
            log_model=True,
        )

    # Build model
    model_config = ModelConfig.load(model_config_path)
    model = Model.from_config(model_config, trust=True)

    # Configuration
    config = TrainingConfig(
        target=[
            target for task in model._tasks for target in task._target_labels
        ],
        early_stopping_patience=early_stopping_patience,
        fit={
            "gpus": gpus,
            "max_epochs": max_epochs,
        },
        dataloader={"batch_size": batch_size, "num_workers": num_workers},
    )

    if suffix is not None:
        archive = os.path.join(EXAMPLE_OUTPUT_DIR, f"train_model_{suffix}")
    else:
        archive = os.path.join(EXAMPLE_OUTPUT_DIR, "train_model")
    run_name = "dynedge_{}_example".format("_".join(config.target))

    # Construct dataloaders
    dataset_config = DatasetConfig.load(dataset_config_path)
    datasets: Dict[str, Any] = Dataset.from_config(
        dataset_config,
    )

    # Construct datasets from multiple selections
    train_dataset = EnsembleDataset(
        [datasets[key] for key in datasets if key.startswith("train")]
    )
    valid_dataset = EnsembleDataset(
        [datasets[key] for key in datasets if key.startswith("valid")]
    )
    test_dataset = EnsembleDataset(
        [datasets[key] for key in datasets if key.startswith("test")]
    )

    # Construct dataloaders
    train_dataloaders = DataLoader(
        train_dataset, shuffle=True, **config.dataloader
    )
    valid_dataloaders = DataLoader(
        valid_dataset, shuffle=False, **config.dataloader
    )
    test_dataloaders = DataLoader(
        test_dataset, shuffle=False, **config.dataloader
    )

    # Log configurations to W&B
    # NB: Only log to W&B on the rank-zero process in case of multi-GPU
    #     training.
    if rank_zero_only == 0:
        wandb_logger.experiment.config.update(config)
        wandb_logger.experiment.config.update(model_config.as_dict())
        wandb_logger.experiment.config.update(dataset_config.as_dict())

    # Build model
    model_config = ModelConfig.load(model_config_path)
    model = Model.from_config(model_config, trust=True)

    # Training model
    callbacks = [
        EarlyStopping(
            monitor="val_loss",
            patience=config.early_stopping_patience,
        ),
        ProgressBar(),
    ]

    model.fit(
        train_dataloaders,
        valid_dataloaders,
        callbacks=callbacks,
        logger=wandb_logger if wandb else None,
        **config.fit,
    )

    # Get predictions
    if isinstance(config.target, str):
        prediction_columns = [
            config.target + "_noise_pred",
            config.target + "_muon_pred",
            config.target + "_neutrino_pred",
        ]
        additional_attributes = [config.target]
    else:
        prediction_columns = [target + "_pred" for target in config.target]
        additional_attributes = config.target

    if prediction_names:
        prediction_columns = prediction_names

    results = model.predict_as_dataframe(
        test_dataloaders,
        prediction_columns=prediction_columns,
        additional_attributes=additional_attributes + ["event_no"],
    )

    # Save predictions and model to file
    db_name = dataset_config.path.split("/")[-1].split(".")[0]
    path = os.path.join(archive, db_name, run_name)
    logger.info(f"Writing results to {path}")
    os.makedirs(path, exist_ok=True)

    results.to_csv(f"{path}/results.csv")
    model.save_state_dict(f"{path}/state_dict.pth")
    model.save(f"{path}/model.pth")


if __name__ == "__main__":
    # Parse command-line arguments
    parser = ArgumentParser(
        description="""
            Train GNN classification model.
            """
    )

    parser.with_standard_arguments(
        (
            "dataset-config",
            os.path.join(
                DATASETS_CONFIG_DIR,
                "training_classification_example_data_sqlite.yml",
            ),
        ),
        (
            "model-config",
            os.path.join(
                MODEL_CONFIG_DIR, "dynedge_PID_classification_example.yml"
            ),
        ),
        "gpus",
        ("max-epochs", 5),
        "early-stopping-patience",
        ("batch-size", 16),
        "num-workers",
    )

    parser.add_argument(
        "--prediction-names",
        nargs="+",
        help="Names of each prediction output feature (default: %(default)s)",
        default=["noise", "muon", "neutrino"],
    )

    parser.add_argument(
        "--suffix",
        type=str,
        help="Name addition to folder (default: %(default)s)",
        default=None,
    )

    args = parser.parse_args()

    main(
        args.dataset_config,
        args.model_config,
        args.gpus,
        args.max_epochs,
        args.early_stopping_patience,
        args.batch_size,
        args.num_workers,
        args.prediction_names,
        args.suffix,
    )
