Note
Go to the end to download the full example code.
Example training script
This script demonstrates how to train a model using the ITI framework. The example shows the translation between PROBA2/SWAP and SDO/AIA.
import os
import collections.abc
import shutil
from itipy.evaluation.util import download_gcp_bucket
#hyper needs the four following aliases to be done manually.
collections.Iterable = collections.abc.Iterable
collections.Mapping = collections.abc.Mapping
collections.MutableSet = collections.abc.MutableSet
collections.MutableMapping = collections.abc.MutableMapping
#Now import hyper
import torch
import yaml
import wandb
from lightning import Trainer
from lightning.pytorch.callbacks import ModelCheckpoint
from lightning.pytorch.loggers import WandbLogger
from sunpy.visualization.colormaps import cm
from itipy.callback import SaveCallback, PlotBAB, PlotABA
from itipy.data.dataset import AIADataset, StorageDataset, SWAPDataset
from itipy.data.data_module import ITIDataModule
from itipy.data.editor import RandomPatchEditor, BrightestPixelPatchEditor
from itipy.iti import ITIModule
Download itipy testset for training
base_path = os.getcwd()
download_gcp_bucket('iti-dataset', base_path+'/iti-testset/')
Failed to download configs/ due to exception: [Errno 21] Is a directory: '/home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/configs/'
Downloaded configs/proba2_to_sdo_colab.yaml to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/configs/proba2_to_sdo_colab.yaml.
Failed to download models/ due to exception: [Errno 21] Is a directory: '/home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/models/'
Downloaded models/aia_to_hri_v0_1.pt to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/models/aia_to_hri_v0_1.pt.
Downloaded models/fsi_to_aia_v0_3.pt to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/models/fsi_to_aia_v0_3.pt.
Downloaded models/swap_to_aia_v0_4.pt to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/models/swap_to_aia_v0_4.pt.
Failed to download proba2/174/ due to exception: [Errno 21] Is a directory: '/home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/proba2/174/'
Downloaded proba2/174/2021-11-11T00:00:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/proba2/174/2021-11-11T00:00:00.fits.
Downloaded proba2/174/2021-11-11T12:00:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/proba2/174/2021-11-11T12:00:00.fits.
Downloaded proba2/174/2021-11-14T00:00:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/proba2/174/2021-11-14T00:00:00.fits.
Downloaded proba2/174/2021-11-15T00:00:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/proba2/174/2021-11-15T00:00:00.fits.
Downloaded proba2/174/2021-11-16T00:00:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/proba2/174/2021-11-16T00:00:00.fits.
Failed to download sdo/ due to exception: [Errno 21] Is a directory: '/home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/sdo/'
Failed to download sdo/171/ due to exception: [Errno 21] Is a directory: '/home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/sdo/171/'
Downloaded sdo/171/2021-11-11T00:00:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/sdo/171/2021-11-11T00:00:00.fits.
Downloaded sdo/171/2021-11-11T12:00:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/sdo/171/2021-11-11T12:00:00.fits.
Downloaded sdo/171/2021-11-14T00:00:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/sdo/171/2021-11-14T00:00:00.fits.
Downloaded sdo/171/2021-11-15T00:00:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/sdo/171/2021-11-15T00:00:00.fits.
Downloaded sdo/171/2021-11-16T00:00:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/sdo/171/2021-11-16T00:00:00.fits.
Failed to download sdo/304/ due to exception: [Errno 21] Is a directory: '/home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/sdo/304/'
Downloaded sdo/304/2021-11-11T00:00:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/sdo/304/2021-11-11T00:00:00.fits.
Downloaded sdo/304/2021-11-11T12:00:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/sdo/304/2021-11-11T12:00:00.fits.
Downloaded sdo/304/2021-11-14T00:00:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/sdo/304/2021-11-14T00:00:00.fits.
Downloaded sdo/304/2021-11-15T00:00:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/sdo/304/2021-11-15T00:00:00.fits.
Downloaded sdo/304/2021-11-16T00:38:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/sdo/304/2021-11-16T00:38:00.fits.
Downloaded soho/171/2007-12-13T07:19:35.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/soho/171/2007-12-13T07:19:35.fits.
Downloaded soho/195/2007-12-13T07:19:35.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/soho/195/2007-12-13T07:19:35.fits.
Downloaded soho/284/2007-12-13T07:19:35.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/soho/284/2007-12-13T07:19:35.fits.
Downloaded soho/304/2007-12-13T07:19:35.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/soho/304/2007-12-13T07:19:35.fits.
Downloaded soho/mag/2007-12-13T07:19:35.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/soho/mag/2007-12-13T07:19:35.fits.
Downloaded solo/.DS_Store to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/solo/.DS_Store.
Downloaded solo/eui-fsi174-image/2021-11-11T00:00:45.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/solo/eui-fsi174-image/2021-11-11T00:00:45.fits.
Downloaded solo/eui-fsi174-image/2021-11-11T12:00:45.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/solo/eui-fsi174-image/2021-11-11T12:00:45.fits.
Downloaded solo/eui-fsi174-image/2021-11-14T00:00:45.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/solo/eui-fsi174-image/2021-11-14T00:00:45.fits.
Downloaded solo/eui-fsi174-image/2021-11-15T00:00:45.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/solo/eui-fsi174-image/2021-11-15T00:00:45.fits.
Downloaded solo/eui-fsi174-image/2021-11-16T00:30:45.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/solo/eui-fsi174-image/2021-11-16T00:30:45.fits.
Failed to download solo/eui-fsi304-image/ due to exception: [Errno 21] Is a directory: '/home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/solo/eui-fsi304-image/'
Downloaded solo/eui-fsi304-image/2021-11-11T00:00:45.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/solo/eui-fsi304-image/2021-11-11T00:00:45.fits.
Downloaded solo/eui-fsi304-image/2021-11-11T12:00:45.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/solo/eui-fsi304-image/2021-11-11T12:00:45.fits.
Downloaded solo/eui-fsi304-image/2021-11-14T00:00:45.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/solo/eui-fsi304-image/2021-11-14T00:00:45.fits.
Downloaded solo/eui-fsi304-image/2021-11-15T00:00:45.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/solo/eui-fsi304-image/2021-11-15T00:00:45.fits.
Downloaded solo/eui-fsi304-image/2021-11-16T00:30:45.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/solo/eui-fsi304-image/2021-11-16T00:30:45.fits.
Downloaded solo/eui-hrieuv174-image/solo_L2_eui-hrieuv174-image_20210220T150905235_V04.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/solo/eui-hrieuv174-image/solo_L2_eui-hrieuv174-image_20210220T150905235_V04.fits.
Downloaded solo/eui-hrieuv174-image/solo_L2_eui-hrieuv174-image_20210220T150910235_V04.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/solo/eui-hrieuv174-image/solo_L2_eui-hrieuv174-image_20210220T150910235_V04.fits.
Downloaded solo/eui-hrieuv174-image/solo_L2_eui-hrieuv174-image_20210220T150915235_V04.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/solo/eui-hrieuv174-image/solo_L2_eui-hrieuv174-image_20210220T150915235_V04.fits.
Downloaded solo/eui-hrieuv174-image/solo_L2_eui-hrieuv174-image_20210220T150920235_V04.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/solo/eui-hrieuv174-image/solo_L2_eui-hrieuv174-image_20210220T150920235_V04.fits.
Downloaded stereo/171/2016-12-28T00:14:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/stereo/171/2016-12-28T00:14:00.fits.
Downloaded stereo/195/2016-12-28T00:14:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/stereo/195/2016-12-28T00:14:00.fits.
Downloaded stereo/284/2016-12-28T00:14:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/stereo/284/2016-12-28T00:14:00.fits.
Downloaded stereo/304/2016-12-28T00:14:00.fits to /home/docs/checkouts/readthedocs.org/user_builds/iti-documentation/checkouts/latest/examples_docs/training/iti-testset/stereo/304/2016-12-28T00:14:00.fits.
Open the config file
with open('iti-testset/configs/proba2_to_sdo_colab.yaml', "r") as stream:
try:
config = yaml.safe_load(stream)
except yaml.YAMLError as exc:
print(exc)
print(config)
base_dir = config['base_dir']
os.makedirs(base_dir, exist_ok=True)
{'base_dir': 'Proba2_to_SDO', 'data': {'B_path': '/iti-testset/proba2/174', 'converted_B_path': 'converted_sdo', 'A_path': '/iti-testset/sdo/171', 'converted_A_path': 'converted_swap', 'num_workers': 4, 'iterations_per_epoch': 10}, 'model': {'input_dim_a': 1, 'input_dim_b': 1, 'upsampling': 1, 'discriminator_mode': 'SINGLE', 'lambda_diversity': 0, 'norm': 'in_rs_aff', 'use_batch_statistic': False}, 'logging': {'wandb_entity': 'christoph-schirninger', 'wandb_project': 'ITI', 'wandb_name': 'PROBA2_to_SDO', 'wandb_id': None}, 'training': {'epochs': 5}}
Initialize the paths from the config file
data_config = config['data']
swap_path = base_path + data_config['A_path']
swap_converted_path = data_config['converted_A_path']
os.makedirs(swap_converted_path, exist_ok=True)
sdo_path = base_path + data_config['B_path']
sdo_converted_path = data_config['converted_B_path']
os.makedirs(sdo_converted_path, exist_ok=True)
Define the training and validation months
Here we use the same training and validation months, as the testset contains only observations from one month. For a real training, you should split the data into training and validation sets (e.g. training=months 2-10, validation=months 11-12).
test_months = list(range(2, 12))
train_months = list(range(2, 12))
Initialize the datasets
To accelerate the training procedure, the data is stored as .npy files in the converted path directories.
The training is performed on patches of size (128, 128) for the SWAP data and (256, 256) for the AIA data. We use the RandomPatchEditor to extract random patches from the full images.
sdo_dataset = AIADataset(sdo_path, wavelength=171, months=train_months)
sdo_dataset = StorageDataset(sdo_dataset,
sdo_converted_path,
ext_editors=[RandomPatchEditor((256, 256))])
swap_dataset = SWAPDataset(swap_path, months=train_months)
swap_dataset = StorageDataset(swap_dataset, swap_converted_path,
ext_editors=[RandomPatchEditor((128, 128))])
sdo_valid = StorageDataset(AIADataset(sdo_path, wavelength=171, months=test_months),
sdo_converted_path, ext_editors=[RandomPatchEditor((256, 256))])
swap_valid = StorageDataset(SWAPDataset(swap_path, months=test_months),
swap_converted_path, ext_editors=[RandomPatchEditor((128, 128))])
Initialize the data module
data_module = ITIDataModule(swap_dataset, sdo_dataset, swap_valid, sdo_valid, **config['data'])
Define the plot settings for the callbacks
plot_settings_A = [
{"cmap": cm.sdoaia171, "title": "SWAP 174", 'vmin': -1, 'vmax': 1},
]
plot_settings_B = [
{"cmap": cm.sdoaia171, "title": "AIA 171", 'vmin': -1, 'vmax': 1},
]
Setup the logging for weights and biases (wandb)
The logging is disabled by default. To enable it, please uncomment the following lines and provide your wandb project and entity name.
logging_config = config[‘logging’] wandb_id = logging_config[‘wandb_id’] if ‘wandb_id’ in logging_config else None log_model = logging_config[‘wandb_log_model’] if ‘wandb_log_model’ in logging_config else False wandb_logger = WandbLogger(project=logging_config[‘wandb_project’], name=logging_config[‘wandb_name’], offline=False,
entity=logging_config[‘wandb_entity’], id=wandb_id, dir=config[‘base_dir’], log_model=log_model)
wandb_logger.experiment.config.update(config, allow_val_change=True)
Start training
module = ITIModule(**config['model'])
Setup save callbacks
checkpoint_callback = ModelCheckpoint(dirpath=base_dir, save_last=True, every_n_epochs=1, save_weights_only=False)
save_callback = SaveCallback(base_dir)
setup plot callbacks
This will plot the predictions of the model on the validation data every epoch.
plot_callbacks = []
#plot_callbacks += [PlotBAB(sdo_valid.sample(2), module, plot_settings_A=plot_settings_A, plot_settings_B=plot_settings_B)]
#plot_callbacks += [PlotABA(swap_valid.sample(2), module, plot_settings_A=plot_settings_A, plot_settings_B=plot_settings_B)]
Set up training parameters
n_gpus = torch.cuda.device_count()
n_cpus = os.cpu_count()
trainer = Trainer(max_epochs=int(config['training']['epochs']),
# #logger=wandb_logger,
devices=n_cpus if n_cpus > 0 else None,
accelerator="cpu" if n_cpus >= 1 else 'auto',
strategy='auto', # ddp breaks memory and wandb
num_sanity_val_steps=0,
callbacks=[checkpoint_callback, save_callback, *plot_callbacks],)
INFO: GPU available: False, used: False
INFO: TPU available: False, using: 0 TPU cores
INFO: HPU available: False, using: 0 HPUs
Start the training
This line is commented out, as the training requires a certain amount of storage space which is not available for a jupyter notebook.
#trainer.fit(module, data_module, ckpt_path='last')
Total running time of the script: (0 minutes 7.379 seconds)