.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "generated/gallery/training/swap_to_aia.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_generated_gallery_training_swap_to_aia.py: ======================= 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. .. GENERATED FROM PYTHON SOURCE LINES 8-36 .. code-block:: Python 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 .. GENERATED FROM PYTHON SOURCE LINES 37-39 Download itipy testset for training .. GENERATED FROM PYTHON SOURCE LINES 39-42 .. code-block:: Python base_path = os.getcwd() download_gcp_bucket('iti-dataset', base_path+'/iti-testset/') .. rst-class:: sphx-glr-script-out .. code-block:: none 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. .. GENERATED FROM PYTHON SOURCE LINES 43-44 Open the config file .. GENERATED FROM PYTHON SOURCE LINES 44-54 .. code-block:: Python 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) .. rst-class:: sphx-glr-script-out .. code-block:: none {'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}} .. GENERATED FROM PYTHON SOURCE LINES 55-56 Initialize the paths from the config file .. GENERATED FROM PYTHON SOURCE LINES 56-64 .. code-block:: Python 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) .. GENERATED FROM PYTHON SOURCE LINES 65-69 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). .. GENERATED FROM PYTHON SOURCE LINES 69-72 .. code-block:: Python test_months = list(range(2, 12)) train_months = list(range(2, 12)) .. GENERATED FROM PYTHON SOURCE LINES 73-78 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. .. GENERATED FROM PYTHON SOURCE LINES 78-94 .. code-block:: Python 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))]) .. GENERATED FROM PYTHON SOURCE LINES 95-96 Initialize the data module .. GENERATED FROM PYTHON SOURCE LINES 96-99 .. code-block:: Python data_module = ITIDataModule(swap_dataset, sdo_dataset, swap_valid, sdo_valid, **config['data']) .. GENERATED FROM PYTHON SOURCE LINES 100-101 Define the plot settings for the callbacks .. GENERATED FROM PYTHON SOURCE LINES 101-109 .. code-block:: Python 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}, ] .. GENERATED FROM PYTHON SOURCE LINES 110-120 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) .. GENERATED FROM PYTHON SOURCE LINES 123-124 Start training .. GENERATED FROM PYTHON SOURCE LINES 124-127 .. code-block:: Python module = ITIModule(**config['model']) .. GENERATED FROM PYTHON SOURCE LINES 128-129 Setup save callbacks .. GENERATED FROM PYTHON SOURCE LINES 129-132 .. code-block:: Python checkpoint_callback = ModelCheckpoint(dirpath=base_dir, save_last=True, every_n_epochs=1, save_weights_only=False) save_callback = SaveCallback(base_dir) .. GENERATED FROM PYTHON SOURCE LINES 133-136 setup plot callbacks This will plot the predictions of the model on the validation data every epoch. .. GENERATED FROM PYTHON SOURCE LINES 136-141 .. code-block:: Python 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)] .. GENERATED FROM PYTHON SOURCE LINES 142-143 Set up training parameters .. GENERATED FROM PYTHON SOURCE LINES 143-155 .. code-block:: Python 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],) .. rst-class:: sphx-glr-script-out .. code-block:: none INFO: GPU available: False, used: False INFO: TPU available: False, using: 0 TPU cores INFO: HPU available: False, using: 0 HPUs .. GENERATED FROM PYTHON SOURCE LINES 156-159 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. .. GENERATED FROM PYTHON SOURCE LINES 159-160 .. code-block:: Python #trainer.fit(module, data_module, ckpt_path='last') .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 7.379 seconds) .. _sphx_glr_download_generated_gallery_training_swap_to_aia.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: swap_to_aia.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: swap_to_aia.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: swap_to_aia.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_