├── Pipfile
├── README.md
├── configs
    ├── baseline.json
    ├── cnn_ae_fd1.json
    ├── cnn_ae_fd2.json
    ├── cnn_ae_fd3.json
    ├── cnn_ae_fd4.json
    ├── cnn_fd1.json
    ├── cnn_fd2.json
    ├── cnn_fd3.json
    ├── cnn_fd4.json
    ├── cnn_pre_fd1.json
    ├── cnn_pre_fd2.json
    ├── cnn_pre_fd3.json
    └── cnn_pre_fd4.json
├── results
    └── semi_supervised.csv
├── src
    ├── building
    │   ├── __init__.py
    │   ├── build.py
    │   └── build_common.py
    ├── datasets
    │   ├── __init__.py
    │   ├── baseline.py
    │   ├── cmapss.py
    │   └── loader.py
    ├── evaluation
    │   ├── add_rul_score.py
    │   ├── export_results_semi_supervised.py
    │   └── plot_semi_supervised.py
    ├── hyperopt
    │   ├── hyperopt_semi_supervised.py
    │   └── hyperopt_supervised.py
    ├── lightning
    │   ├── autoencoder.py
    │   ├── baseline.py
    │   ├── loggers.py
    │   ├── metrics.py
    │   ├── mixins.py
    │   └── pretraining.py
    ├── models
    │   ├── layers.py
    │   └── networks.py
    ├── reproduce_baseline.py
    ├── reproduce_semi_supervised.py
    ├── run_baseline.py
    ├── run_pretraining.py
    ├── run_rbm.py
    └── run_semi_supervised.py
└── tests
    ├── __init__.py
    ├── dataset_tests
        ├── __init__.py
        ├── templates.py
        ├── test_baseline.py
        ├── test_cmapss.py
        └── test_loader.py
    ├── lightning_tests
        ├── __init__.py
        ├── test_autoencoder.py
        ├── test_baseline.py
        ├── test_loggers.py
        ├── test_metrics.py
        └── test_pretraining.py
    ├── model_tests
        ├── __init__.py
        ├── test_layers.py
        └── test_networks.py
    └── test_building.py


/Pipfile:
--------------------------------------------------------------------------------
 1 | [[source]]
 2 | url = "https://pypi.org/simple"
 3 | verify_ssl = true
 4 | name = "pypi"
 5 | 
 6 | [[source]]
 7 | name = "pytorch"
 8 | url = "https://download.pytorch.org/whl/cu102"
 9 | verify_ssl = false
10 | 
11 | [packages]
12 | plotnine = "==0.7.1"
13 | numpy = "==1.19.1"
14 | mlflow = "==1.10.0"
15 | pytorch-lightning = "==0.10.0"
16 | pandas = "==1.1.3"
17 | tabulate = "==0.8.7"
18 | matplotlib = "*"
19 | tqdm = "==4.50.2"
20 | ray = "==1.0.0"
21 | scikit-learn = "==0.23.2"
22 | umap-learn = "==0.4.6"
23 | numba = "==0.51.2"
24 | PyTorch-ProbGraph = "==0.0.1"
25 | torch = "==1.6.0"
26 | 
27 | [dev-packages]
28 | 
29 | [requires]
30 | python_version = "3.7"
31 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # Improving Semi-Supervised Learning for Remaining Useful Lifetime Estimation Through Self-Supervision
 2 | 
 3 | This is the companion repository for the paper "*Improving Semi-Supervised Learning for Remaining Useful Lifetime Estimation Through Self-Supervision*".
 4 | It is available as a manuscript on [arXiv](https://arxiv.org/abs/2108.08721) or as a full journal article in [IJPHM](https://papers.phmsociety.org/index.php/ijphm/article/view/3096).
 5 | 
 6 | ## Installation
 7 | 
 8 | ### Environment
 9 | This project uses Python 3.7 and pipenv to manage its environment. Install pipenv as:
10 | 
11 | ```
12 | pip install -U pipenv
13 | ```
14 | 
15 | You can set up this project in a virtual environment afterwards with:
16 | 
17 | ```shell
18 | cd <PROJECT_ROOT>
19 | pipenv install --skip-lock
20 | ```
21 | 
22 | Please make sure that the appropriate Nvidia driver, CUDA 10.2 and cuDNN are installed on your system.
23 | 
24 | ### Data
25 | 
26 | Download the NASA C-MAPSS dataset [here](https://ti.arc.nasa.gov/c/6/) and extract the content to ``<PROJECT_ROOT>/data/CMAPSS``.
27 | 
28 | ### Verify Installation
29 | 
30 | Test the installation by running the unit tests:
31 | 
32 | ```shell
33 | pipenv shell
34 | export PYTHONPATH=$PYTHONPATH:./src
35 | python -m unittest -v
36 | ```
37 | 
38 | ## Usage
39 | 
40 | ### Concerning Ray
41 | 
42 | This project uses ``ray`` to parallelize the training.
43 | It expects a GPU and at least 3 CPU cores to run one training instance and will try to put two runs on one GPU if more CPU cores are available.
44 | If you want to customize these requirements, please edit the ``ray.remote`` function decorators in the file ``run_semi_supervised.py``.
45 | 
46 | ### Basic Usage
47 | 
48 | To start reproducing the results, enter the source directory of the project and open the pipenv shell if you haven't already:
49 | 
50 | ```shell
51 | pipenv shell
52 | cd <PROJECT_ROOT>/src
53 | ```
54 | 
55 | To reproduce all baseline experiments call:
56 | 
57 | ```shell
58 | python reproduce_baseline.py 2021
59 | ```
60 | 
61 | To reproduce each of the SSL methods call:
62 | 
63 | ```shell
64 | python reproduce_semi_supervised.py 2021 --mode metric
65 | python reproduce_semi_supervised.py autoencoder --mode autoencoder
66 | python reproduce_semi_supervised.py rbm --mode rbm
67 | ```
68 | 
69 | ### Advanced Usage
70 | 
71 | If you want to modify any of the experiments, you can start with the hyperparameter configurations.
72 | The ``config`` folder contains JSON configs for each experiment.
73 | 
74 | If you want to run an autoencoder or self-supervised pre-training experiment, you can use the ``run_pretraining.py`` script.
75 | Pre-training experiments with RBMs can be run with the ``run_rbm.py`` script.
76 | Please consult the scripts help page for their usage with ``python <SCRIPT> --help``.
77 | 


--------------------------------------------------------------------------------
/configs/baseline.json:
--------------------------------------------------------------------------------
1 | {
2 |     "num_layers": 6,
3 |     "base_filters": 16,
4 |     "latent_dim": 64,
5 |     "dropout": 0.1,
6 |     "lr": 0.01,
7 |     "batch_size": 256,
8 |     "num_disc_layers": 1
9 | }


--------------------------------------------------------------------------------
/configs/cnn_ae_fd1.json:
--------------------------------------------------------------------------------
1 | {
2 |     "domain_tradeoff": 0.0,
3 |     "dropout": 0.1,
4 |     "lr": 0.0001,
5 |     "batch_size": 64,
6 |     "distance_mode": "linear",
7 |     "min_distance": 1
8 | }


--------------------------------------------------------------------------------
/configs/cnn_ae_fd2.json:
--------------------------------------------------------------------------------
1 | {
2 |     "domain_tradeoff": 0.0,
3 |     "dropout": 0.4,
4 |     "lr": 0.0248,
5 |     "batch_size": 256,
6 |     "distance_mode": "linear",
7 |     "min_distance": 15
8 | }


--------------------------------------------------------------------------------
/configs/cnn_ae_fd3.json:
--------------------------------------------------------------------------------
1 | {
2 |     "domain_tradeoff": 0.0,
3 |     "dropout": 0.0,
4 |     "lr": 0.015,
5 |     "batch_size": 64,
6 |     "distance_mode": "linear",
7 |     "min_distance": 1
8 | }


--------------------------------------------------------------------------------
/configs/cnn_ae_fd4.json:
--------------------------------------------------------------------------------
1 | {
2 |     "domain_tradeoff": 0.0,
3 |     "dropout": 0.0,
4 |     "lr": 0.0006,
5 |     "batch_size": 64,
6 |     "distance_mode": "linear",
7 |     "min_distance": 10
8 | }


--------------------------------------------------------------------------------
/configs/cnn_fd1.json:
--------------------------------------------------------------------------------
1 | {
2 |     "num_layers": 6,
3 |     "base_filters": 16,
4 |     "latent_dim": 64,
5 |     "dropout": 0.4,
6 |     "lr": 0.0056,
7 |     "batch_size": 128,
8 |     "num_disc_layers": 1
9 | }


--------------------------------------------------------------------------------
/configs/cnn_fd2.json:
--------------------------------------------------------------------------------
1 | {
2 |     "num_layers": 6,
3 |     "base_filters": 16,
4 |     "latent_dim": 64,
5 |     "dropout": 0.3,
6 |     "lr": 0.0903,
7 |     "batch_size": 512,
8 |     "num_disc_layers": 1
9 | }


--------------------------------------------------------------------------------
/configs/cnn_fd3.json:
--------------------------------------------------------------------------------
1 | {
2 |     "num_layers": 6,
3 |     "base_filters": 16,
4 |     "latent_dim": 64,
5 |     "dropout": 0.2,
6 |     "lr": 0.095,
7 |     "batch_size": 64,
8 |     "num_disc_layers": 1
9 | }


--------------------------------------------------------------------------------
/configs/cnn_fd4.json:
--------------------------------------------------------------------------------
1 | {
2 |     "num_layers": 6,
3 |     "base_filters": 16,
4 |     "latent_dim": 64,
5 |     "dropout": 0.0,
6 |     "lr": 0.06635,
7 |     "batch_size": 64,
8 |     "num_disc_layers": 1
9 | }


--------------------------------------------------------------------------------
/configs/cnn_pre_fd1.json:
--------------------------------------------------------------------------------
1 | {
2 |     "domain_tradeoff": 0.0,
3 |     "dropout": 0.2,
4 |     "lr": 0.00015,
5 |     "batch_size": 64,
6 |     "distance_mode": "linear",
7 |     "min_distance": 10
8 | }


--------------------------------------------------------------------------------
/configs/cnn_pre_fd2.json:
--------------------------------------------------------------------------------
1 | {
2 |     "domain_tradeoff": 0.0,
3 |     "dropout": 0.4,
4 |     "lr": 0.01155,
5 |     "batch_size": 64,
6 |     "distance_mode": "linear",
7 |     "min_distance": 15
8 | }


--------------------------------------------------------------------------------
/configs/cnn_pre_fd3.json:
--------------------------------------------------------------------------------
1 | {
2 |     "domain_tradeoff": 0.0,
3 |     "dropout": 0.1,
4 |     "lr": 0.00615,
5 |     "batch_size": 64,
6 |     "distance_mode": "linear",
7 |     "min_distance": 15
8 | }


--------------------------------------------------------------------------------
/configs/cnn_pre_fd4.json:
--------------------------------------------------------------------------------
1 | {
2 |     "domain_tradeoff": 0.0,
3 |     "dropout": 0.1,
4 |     "lr": 0.07455,
5 |     "batch_size": 64,
6 |     "distance_mode": "linear",
7 |     "min_distance": 10
8 | }


--------------------------------------------------------------------------------
/src/building/__init__.py:
--------------------------------------------------------------------------------
 1 | from .build import (
 2 |     build_autoencoder_from_config,
 3 |     build_baseline,
 4 |     build_baseline_from_config,
 5 |     build_datamodule,
 6 |     build_pretraining,
 7 |     build_pretraining_from_config,
 8 |     build_rbm,
 9 | )
10 | from .build_common import add_hparams, build_trainer, get_logdir, load_config
11 | 


--------------------------------------------------------------------------------
/src/building/build.py:
--------------------------------------------------------------------------------
  1 | import pytorch_lightning as pl
  2 | import torch
  3 | from pytorch_probgraph import (
  4 |     InteractionModule,
  5 |     RestrictedBoltzmannMachineCD,
  6 | )
  7 | 
  8 | import datasets
  9 | from building.build_common import add_hparams, build_trainer, get_logdir
 10 | from lightning import autoencoder, baseline, loggers, pretraining
 11 | from models.layers import GaussianSequenceLayer, RectifiedLinearLayer
 12 | 
 13 | 
 14 | def build_baseline(
 15 |     source,
 16 |     fails,
 17 |     config,
 18 |     encoder,
 19 |     pretrained_encoder_path,
 20 |     gpu,
 21 |     seed,
 22 |     version,
 23 |     record_embeddings=False,
 24 | ):
 25 |     logger = loggers.MLTBLogger(
 26 |         get_logdir(), loggers.baseline_experiment_name(source), tag=version
 27 |     )
 28 |     checkpoint_callback = pl.callbacks.ModelCheckpoint(monitor="val/regression_loss")
 29 |     trainer = build_trainer(
 30 |         logger,
 31 |         checkpoint_callback,
 32 |         max_epochs=100,
 33 |         val_interval=1.0,
 34 |         gpu=gpu,
 35 |         seed=seed,
 36 |     )
 37 |     data = datasets.BaselineDataModule(
 38 |         fd_source=source, batch_size=config["batch_size"], percent_fail_runs=fails
 39 |     )
 40 |     model = build_baseline_from_config(
 41 |         config, data.window_size, encoder, pretrained_encoder_path, record_embeddings
 42 |     )
 43 |     add_hparams(model, data, seed)
 44 | 
 45 |     return trainer, data, model
 46 | 
 47 | 
 48 | def build_baseline_from_config(
 49 |     config, seq_len, encoder, pretrained_encoder_path, record_embeddings=False
 50 | ):
 51 |     model = baseline.Baseline(
 52 |         in_channels=14,
 53 |         seq_len=seq_len,
 54 |         num_layers=config["num_layers"],
 55 |         kernel_size=3,
 56 |         base_filters=config["base_filters"],
 57 |         latent_dim=config["latent_dim"],
 58 |         dropout=config["dropout"],
 59 |         optim_type="adam",
 60 |         lr=config["lr"],
 61 |         record_embeddings=record_embeddings,
 62 |         encoder=encoder,
 63 |     )
 64 |     if pretrained_encoder_path is not None:
 65 |         model.load_encoder(pretrained_encoder_path)
 66 | 
 67 |     return model
 68 | 
 69 | 
 70 | def build_pretraining(
 71 |     source,
 72 |     target,
 73 |     percent_broken,
 74 |     percent_fail_runs,
 75 |     arch_config,
 76 |     config,
 77 |     encoder,
 78 |     mode,
 79 |     record_embeddings,
 80 |     gpu,
 81 |     seed,
 82 |     version,
 83 | ):
 84 |     logger = loggers.MLTBLogger(
 85 |         get_logdir(),
 86 |         loggers.pretraining_experiment_name(source, target),
 87 |         tag=version,
 88 |         tensorboard_struct={"pb": percent_broken, "dt": config["domain_tradeoff"]},
 89 |     )
 90 |     checkpoint_callback = loggers.MinEpochModelCheckpoint(
 91 |         monitor="val/checkpoint_score", min_epochs_before_saving=1
 92 |     )
 93 |     trainer = build_trainer(
 94 |         logger,
 95 |         checkpoint_callback,
 96 |         max_epochs=100,
 97 |         val_interval=1.0,
 98 |         gpu=gpu,
 99 |         seed=seed,
100 |     )
101 |     truncate_val = not record_embeddings
102 |     distance_mode = config["distance_mode"]
103 |     min_distance = config["min_distance"] if "min_distance" in config else 1
104 |     data = build_datamodule(
105 |         source,
106 |         target,
107 |         percent_broken,
108 |         percent_fail_runs,
109 |         config["batch_size"],
110 |         truncate_val,
111 |         distance_mode,
112 |         min_distance,
113 |     )
114 |     use_adaption = target is not None
115 |     if mode == "metric":
116 |         model = build_pretraining_from_config(
117 |             arch_config,
118 |             config,
119 |             data.window_size,
120 |             encoder,
121 |             record_embeddings,
122 |             use_adaption,
123 |         )
124 |     elif mode == "autoencoder":
125 |         model = build_autoencoder_from_config(
126 |             arch_config,
127 |             config,
128 |             data.window_size,
129 |             encoder,
130 |             record_embeddings,
131 |             use_adaption,
132 |         )
133 |     else:
134 |         raise ValueError(f"Unrecognized pre-training mode {mode}.")
135 | 
136 |     add_hparams(model, data, seed)
137 | 
138 |     return trainer, data, model
139 | 
140 | 
141 | def build_datamodule(
142 |     source,
143 |     target,
144 |     percent_broken,
145 |     percent_fail_runs,
146 |     batch_size,
147 |     truncate_val,
148 |     distance_mode,
149 |     min_distance,
150 |     window_size=None,
151 | ):
152 |     if target is None:
153 |         return datasets.PretrainingBaselineDataModule(
154 |             fd_source=source,
155 |             num_samples=25000,
156 |             batch_size=batch_size,
157 |             window_size=window_size,
158 |             min_distance=min_distance,
159 |             percent_broken=percent_broken,
160 |             percent_fail_runs=percent_fail_runs,
161 |             truncate_val=truncate_val,
162 |             distance_mode=distance_mode,
163 |         )
164 |     else:
165 |         return datasets.PretrainingAdaptionDataModule(
166 |             fd_source=source,
167 |             fd_target=target,
168 |             num_samples=50000,
169 |             batch_size=batch_size,
170 |             window_size=window_size,
171 |             min_distance=min_distance,
172 |             percent_broken=percent_broken,
173 |             percent_fail_runs=percent_fail_runs,
174 |             truncate_target_val=truncate_val,
175 |             distance_mode=distance_mode,
176 |         )
177 | 
178 | 
179 | def build_pretraining_from_config(
180 |     arch_config, config, seq_len, encoder, record_embeddings, use_adaption
181 | ):
182 |     model = pretraining.UnsupervisedPretraining(
183 |         in_channels=14,
184 |         seq_len=seq_len,
185 |         num_layers=arch_config["num_layers"],
186 |         kernel_size=3,
187 |         base_filters=arch_config["base_filters"],
188 |         latent_dim=arch_config["latent_dim"],
189 |         dropout=config["dropout"],
190 |         domain_tradeoff=config["domain_tradeoff"] if use_adaption else 0.0,
191 |         domain_disc_dim=arch_config["latent_dim"],
192 |         num_disc_layers=arch_config["num_disc_layers"],
193 |         lr=config["lr"],
194 |         weight_decay=0.0,
195 |         record_embeddings=record_embeddings,
196 |         encoder=encoder,
197 |     )
198 | 
199 |     return model
200 | 
201 | 
202 | def build_autoencoder_from_config(
203 |     arch_config, config, seq_len, encoder, record_embeddings, use_adaption
204 | ):
205 |     model = autoencoder.AutoencoderPretraining(
206 |         in_channels=14,
207 |         seq_len=seq_len,
208 |         num_layers=arch_config["num_layers"],
209 |         kernel_size=3,
210 |         base_filters=arch_config["base_filters"],
211 |         latent_dim=arch_config["latent_dim"],
212 |         dropout=config["dropout"],
213 |         domain_tradeoff=config["domain_tradeoff"] if use_adaption else 0.0,
214 |         domain_disc_dim=arch_config["latent_dim"],
215 |         num_disc_layers=arch_config["num_disc_layers"],
216 |         lr=config["lr"],
217 |         weight_decay=0.0,
218 |         record_embeddings=record_embeddings,
219 |         encoder=encoder,
220 |     )
221 | 
222 |     return model
223 | 
224 | 
225 | def build_rbm(in_units, out_units, he_init=False):
226 |     l0bias = torch.zeros([1, in_units, 1])
227 |     l0bias.requires_grad = True
228 |     l1bias = torch.zeros([1, out_units, 1])
229 |     l1bias.requires_grad = True
230 | 
231 |     l0 = GaussianSequenceLayer(l0bias, torch.ones_like(l0bias))
232 |     l0.logsigma.requires_grad = False
233 |     l1 = RectifiedLinearLayer(l1bias)
234 | 
235 |     module = torch.nn.Conv1d(14, out_units, kernel_size=3, bias=False)
236 |     i0 = InteractionModule(module, l0.bias.shape[1:])
237 |     if he_init:
238 |         torch.nn.init.kaiming_uniform_(i0.weight, nonlinearity="relu")
239 | 
240 |     # build the RBM
241 |     rbm = RestrictedBoltzmannMachineCD(l0, l1, i0, ksteps=1)
242 | 
243 |     return rbm
244 | 


--------------------------------------------------------------------------------
/src/building/build_common.py:
--------------------------------------------------------------------------------
 1 | import json
 2 | import os
 3 | 
 4 | import pytorch_lightning as pl
 5 | 
 6 | 
 7 | def build_trainer(
 8 |     logger,
 9 |     checkpoint_callback,
10 |     max_epochs,
11 |     val_interval,
12 |     gpu,
13 |     seed=None,
14 |     callbacks=None,
15 |     check_sanity=True,
16 | ):
17 |     if seed is not None:
18 |         pl.trainer.seed_everything(seed)
19 |     trainer = pl.Trainer(
20 |         num_sanity_val_steps=2 if check_sanity else 0,
21 |         gpus=[gpu],
22 |         max_epochs=max_epochs,
23 |         logger=logger,
24 |         deterministic=True,
25 |         log_every_n_steps=10,
26 |         checkpoint_callback=checkpoint_callback,
27 |         gradient_clip_val=1.0,
28 |         val_check_interval=val_interval,
29 |         callbacks=callbacks,
30 |     )
31 | 
32 |     return trainer
33 | 
34 | 
35 | def add_hparams(model, data, seed):
36 |     model.add_data_hparams(data)
37 |     model.hparams.update({"seed": seed})
38 | 
39 | 
40 | def get_logdir():
41 |     script_path = os.path.dirname(__file__)
42 |     log_dir = os.path.normpath(os.path.join(script_path, "..", ".."))
43 | 
44 |     return log_dir
45 | 
46 | 
47 | def load_config(config_path):
48 |     with open(config_path, mode="rt") as f:
49 |         config = json.load(f)
50 | 
51 |     return config
52 | 


--------------------------------------------------------------------------------
/src/datasets/__init__.py:
--------------------------------------------------------------------------------
1 | from .baseline import BaselineDataModule, PretrainingBaselineDataModule
2 | 


--------------------------------------------------------------------------------
/src/datasets/baseline.py:
--------------------------------------------------------------------------------
  1 | from typing import List, Optional, Union
  2 | 
  3 | import pytorch_lightning as pl
  4 | from torch.utils.data import DataLoader
  5 | 
  6 | from datasets.cmapss import CMAPSSDataModule, PairedCMAPSS
  7 | from datasets.loader import CMAPSSLoader
  8 | 
  9 | 
 10 | class BaselineDataModule(pl.LightningDataModule):
 11 |     def __init__(
 12 |         self,
 13 |         fd_source: int,
 14 |         batch_size: int,
 15 |         window_size: int = None,
 16 |         max_rul: int = 125,
 17 |         percent_fail_runs: float = None,
 18 |         feature_select: List[int] = None,
 19 |     ):
 20 |         super().__init__()
 21 | 
 22 |         self.fd_source = fd_source
 23 |         self.batch_size = batch_size
 24 |         self.max_rul = max_rul
 25 |         self.window_size = window_size or CMAPSSLoader.WINDOW_SIZES[self.fd_source]
 26 |         self.percent_fail_runs = percent_fail_runs
 27 |         self.feature_select = feature_select
 28 | 
 29 |         self.hparams = {
 30 |             "fd_source": self.fd_source,
 31 |             "batch_size": self.batch_size,
 32 |             "window_size": self.window_size,
 33 |             "percent_fail_runs": self.percent_fail_runs,
 34 |             "max_rul": self.max_rul,
 35 |         }
 36 | 
 37 |         self.cmapss = {}
 38 |         for fd in range(1, 5):
 39 |             self.cmapss[fd] = self._get_cmapss(fd)
 40 | 
 41 |     def _get_cmapss(self, fd):
 42 |         fail_runs = self.percent_fail_runs if fd == self.fd_source else None
 43 |         cmapss = CMAPSSDataModule(
 44 |             fd,
 45 |             self.batch_size,
 46 |             self.window_size,
 47 |             self.max_rul,
 48 |             None,
 49 |             fail_runs,
 50 |             self.feature_select,
 51 |         )
 52 | 
 53 |         return cmapss
 54 | 
 55 |     def prepare_data(self, *args, **kwargs):
 56 |         for cmapss_fd in self.cmapss.values():
 57 |             cmapss_fd.prepare_data(*args, **kwargs)
 58 | 
 59 |     def setup(self, stage: Optional[str] = None):
 60 |         for cmapss_fd in self.cmapss.values():
 61 |             cmapss_fd.setup(stage)
 62 | 
 63 |     def train_dataloader(self, *args, **kwargs) -> DataLoader:
 64 |         return self.cmapss[self.fd_source].train_dataloader()
 65 | 
 66 |     def val_dataloader(self, *args, **kwargs) -> DataLoader:
 67 |         return self.cmapss[self.fd_source].val_dataloader()
 68 | 
 69 |     def test_dataloader(self, *args, **kwargs) -> List[DataLoader]:
 70 |         test_dataloaders = []
 71 |         for fd_target in range(1, 5):
 72 |             target_dl = self.cmapss[fd_target].test_dataloader()
 73 |             test_dataloaders.append(target_dl)
 74 | 
 75 |         return test_dataloaders
 76 | 
 77 | 
 78 | class PretrainingBaselineDataModule(pl.LightningDataModule):
 79 |     def __init__(
 80 |         self,
 81 |         fd_source: int,
 82 |         num_samples: int,
 83 |         batch_size: int,
 84 |         window_size: int = None,
 85 |         max_rul: int = 125,
 86 |         min_distance: int = 1,
 87 |         percent_broken: float = None,
 88 |         percent_fail_runs: Union[float, List[int]] = None,
 89 |         feature_select: List[int] = None,
 90 |         truncate_val: bool = False,
 91 |         distance_mode: str = "linear",
 92 |     ):
 93 |         super().__init__()
 94 | 
 95 |         self.fd_source = fd_source
 96 |         self.num_samples = num_samples
 97 |         self.batch_size = batch_size
 98 |         self.min_distance = min_distance
 99 |         self.max_rul = max_rul
100 |         self.percent_broken = percent_broken
101 |         self.percent_fail_runs = percent_fail_runs
102 |         self.feature_select = feature_select
103 |         self.truncate_val = truncate_val
104 |         self.distance_mode = distance_mode
105 | 
106 |         self.fails_source_loader = CMAPSSLoader(
107 |             self.fd_source,
108 |             window_size,
109 |             self.max_rul,
110 |             None,
111 |             self.percent_fail_runs,
112 |             self.feature_select,
113 |             self.truncate_val,
114 |         )
115 |         unfailed_runs = self._get_unbroken_runs(self.percent_fail_runs)
116 |         self.broken_source_loader = CMAPSSLoader(
117 |             self.fd_source,
118 |             window_size,
119 |             self.max_rul,
120 |             self.percent_broken,
121 |             unfailed_runs,
122 |             self.feature_select,
123 |             self.truncate_val,
124 |         )
125 | 
126 |         self.window_size = self.broken_source_loader.window_size
127 | 
128 |         self.source = CMAPSSDataModule.from_loader(
129 |             self.broken_source_loader, self.batch_size
130 |         )
131 | 
132 |         self.hparams = {
133 |             "fd_source": self.fd_source,
134 |             "num_samples": self.num_samples,
135 |             "batch_size": self.batch_size,
136 |             "window_size": self.window_size,
137 |             "max_rul": self.max_rul,
138 |             "min_distance": self.min_distance,
139 |             "percent_broken": self.percent_broken,
140 |             "percent_fail_runs": self.percent_fail_runs,
141 |             "truncate_val": self.truncate_val,
142 |             "distance_mode": self.distance_mode,
143 |         }
144 | 
145 |     def _get_unbroken_runs(self, fail_runs):
146 |         if fail_runs is None or isinstance(fail_runs, float):
147 |             unfailed_runs = None
148 |         else:
149 |             run_idx = range(CMAPSSLoader.NUM_TRAIN_RUNS[self.fd_source])
150 |             unfailed_runs = list(set(run_idx).difference(fail_runs))
151 | 
152 |         return unfailed_runs
153 | 
154 |     def prepare_data(self, *args, **kwargs):
155 |         self.broken_source_loader.prepare_data()
156 | 
157 |     def setup(self, stage: Optional[str] = None):
158 |         self.source.setup(stage)
159 | 
160 |     def train_dataloader(self, *args, **kwargs) -> DataLoader:
161 |         return DataLoader(
162 |             self._get_paired_dataset("dev"), batch_size=self.batch_size, pin_memory=True
163 |         )
164 | 
165 |     def val_dataloader(self, *args, **kwargs) -> List[DataLoader]:
166 |         combined_loader = DataLoader(
167 |             self._get_paired_dataset("val"), batch_size=self.batch_size, pin_memory=True
168 |         )
169 |         source_loader = self.source.val_dataloader()
170 | 
171 |         return [combined_loader, source_loader]
172 | 
173 |     def _get_paired_dataset(self, split: str) -> PairedCMAPSS:
174 |         deterministic = split == "val"
175 |         min_distance = 1 if split == "val" else self.min_distance
176 |         num_samples = 25000 if split == "val" else self.num_samples
177 |         paired = PairedCMAPSS(
178 |             [self.broken_source_loader, self.fails_source_loader],
179 |             split,
180 |             num_samples,
181 |             min_distance,
182 |             deterministic,
183 |             mode=self.distance_mode,
184 |         )
185 | 
186 |         return paired
187 | 


--------------------------------------------------------------------------------
/src/datasets/cmapss.py:
--------------------------------------------------------------------------------
  1 | from typing import List, Optional, Tuple, Union
  2 | 
  3 | import numpy as np
  4 | import pytorch_lightning as pl
  5 | import torch
  6 | from torch.utils.data import DataLoader, Dataset, IterableDataset, TensorDataset
  7 | 
  8 | from datasets.loader import CMAPSSLoader
  9 | 
 10 | 
 11 | class CMAPSSDataModule(pl.LightningDataModule):
 12 |     def __init__(
 13 |         self,
 14 |         fd: int,
 15 |         batch_size: int,
 16 |         window_size: int = None,
 17 |         max_rul: int = 125,
 18 |         percent_broken: float = None,
 19 |         percent_fail_runs: Union[float, List[int]] = None,
 20 |         feature_select: int = None,
 21 |         truncate_val: bool = False,
 22 |     ):
 23 |         super().__init__()
 24 | 
 25 |         self._loader = CMAPSSLoader(
 26 |             fd,
 27 |             window_size,
 28 |             max_rul,
 29 |             percent_broken,
 30 |             percent_fail_runs,
 31 |             feature_select,
 32 |             truncate_val,
 33 |         )
 34 | 
 35 |         self.batch_size = batch_size
 36 |         self.fd = self._loader.fd
 37 |         self.window_size = self._loader.window_size
 38 |         self.max_rul = self._loader.max_rul
 39 |         self.percent_broken = self._loader.percent_broken
 40 |         self.percent_fail_runs = self._loader.percent_fail_runs
 41 |         self.feature_select = self._loader.feature_select
 42 |         self.truncate_val = self._loader.truncate_val
 43 | 
 44 |         self.hparams = {
 45 |             "fd": self.fd,
 46 |             "batch_size": self.batch_size,
 47 |             "window_size": self.window_size,
 48 |             "max_rul": self.max_rul,
 49 |             "percent_broken": self.percent_broken,
 50 |             "percent_fail_runs": self.percent_fail_runs,
 51 |             "truncate_val": self.truncate_val,
 52 |         }
 53 | 
 54 |         self.data = {}
 55 | 
 56 |     @classmethod
 57 |     def from_loader(cls, loader: CMAPSSLoader, batch_size: int):
 58 |         return cls(
 59 |             loader.fd,
 60 |             batch_size,
 61 |             loader.window_size,
 62 |             loader.max_rul,
 63 |             loader.percent_broken,
 64 |             loader.percent_fail_runs,
 65 |             loader.feature_select,
 66 |             loader.truncate_val,
 67 |         )
 68 | 
 69 |     def prepare_data(self, *args, **kwargs):
 70 |         self._loader.prepare_data()
 71 | 
 72 |     def setup(self, stage: Optional[str] = None):
 73 |         self.data["dev"] = self._setup_split("dev")
 74 |         self.data["val"] = self._setup_split("val")
 75 |         self.data["test"] = self._setup_split("test")
 76 | 
 77 |     def _setup_split(self, split):
 78 |         features, targets = self._loader.load_split(split)
 79 |         if features:
 80 |             features = torch.cat(features)
 81 |             targets = torch.cat(targets)
 82 |         else:
 83 |             features = torch.empty(0, len(self._loader.feature_select), self.window_size)
 84 |             targets = torch.empty(0)
 85 | 
 86 |         return features, targets
 87 | 
 88 |     def train_dataloader(self, *args, **kwargs) -> DataLoader:
 89 |         dataset = self.to_dataset("dev")
 90 |         drop_last = len(dataset) % self.batch_size == 1
 91 |         loader = DataLoader(
 92 |             dataset,
 93 |             batch_size=self.batch_size,
 94 |             shuffle=True,
 95 |             drop_last=drop_last,
 96 |             pin_memory=True,
 97 |         )
 98 | 
 99 |         return loader
100 | 
101 |     def val_dataloader(self, *args, **kwargs) -> Union[DataLoader, List[DataLoader]]:
102 |         return DataLoader(
103 |             self.to_dataset("val"),
104 |             batch_size=self.batch_size,
105 |             shuffle=False,
106 |             pin_memory=True,
107 |         )
108 | 
109 |     def test_dataloader(self, *args, **kwargs) -> Union[DataLoader, List[DataLoader]]:
110 |         return DataLoader(
111 |             self.to_dataset("test"),
112 |             batch_size=self.batch_size,
113 |             shuffle=False,
114 |             pin_memory=True,
115 |         )
116 | 
117 |     def to_dataset(self, split):
118 |         features, targets = self.data[split]
119 |         split_dataset = TensorDataset(features, targets)
120 | 
121 |         return split_dataset
122 | 
123 | 
124 | class PairedCMAPSS(IterableDataset):
125 |     def __init__(
126 |         self,
127 |         loaders: List[CMAPSSLoader],
128 |         split: str,
129 |         num_samples: int,
130 |         min_distance: int,
131 |         deterministic: bool = False,
132 |         mode: str = "linear",
133 |     ):
134 |         super().__init__()
135 | 
136 |         self.loaders = loaders
137 |         self.split = split
138 |         self.min_distance = min_distance
139 |         self.num_samples = num_samples
140 |         self.deterministic = deterministic
141 |         self.mode = mode
142 | 
143 |         if not all(d.window_size == self.loaders[0].window_size for d in self.loaders):
144 |             window_sizes = [d.window_size for d in self.loaders]
145 |             raise ValueError(
146 |                 f"Datasets to be paired do not have the same window size, but {window_sizes}"
147 |             )
148 | 
149 |         self._run_domain_idx = None
150 |         self._features = None
151 |         self._labels = None
152 |         self._prepare_datasets()
153 | 
154 |         self._max_rul = max(loader.max_rul for loader in self.loaders)
155 |         self._current_iteration = 0
156 |         self._rng = self._reset_rng()
157 |         if mode == "linear":
158 |             self._get_pair_func = self._get_pair_idx
159 |         elif mode == "piecewise":
160 |             self._get_pair_func = self._get_pair_idx_piecewise
161 |         elif mode == "labeled":
162 |             self._get_pair_func = self._get_labeled_pair_idx
163 | 
164 |     def _prepare_datasets(self):
165 |         run_domain_idx = []
166 |         features = []
167 |         labels = []
168 |         for domain_idx, loader in enumerate(self.loaders):
169 |             run_features, run_labels = loader.load_split(self.split)
170 |             for feat, lab in zip(run_features, run_labels):
171 |                 if len(feat) > self.min_distance:
172 |                     run_domain_idx.append(domain_idx)
173 |                     features.append(feat)
174 |                     labels.append(lab)
175 | 
176 |         self._run_domain_idx = np.array(run_domain_idx)
177 |         self._features = features
178 |         self._labels = labels
179 | 
180 |     def _reset_rng(self) -> np.random.Generator:
181 |         return np.random.default_rng(seed=42)
182 | 
183 |     def __len__(self) -> int:
184 |         return self.num_samples
185 | 
186 |     def __iter__(self):
187 |         self._current_iteration = 0
188 |         if self.deterministic:
189 |             self._rng = self._reset_rng()
190 | 
191 |         return self
192 | 
193 |     def __next__(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
194 |         if self._current_iteration < self.num_samples:
195 |             self._current_iteration += 1
196 |             idx = self._get_pair_func()
197 |             return self._build_pair(*idx)
198 |         else:
199 |             raise StopIteration
200 | 
201 |     def _get_pair_idx(self) -> Tuple[torch.Tensor, int, int, int, int]:
202 |         chosen_run_idx = self._rng.integers(0, len(self._features))
203 |         domain_label = self._run_domain_idx[chosen_run_idx]
204 |         chosen_run = self._features[chosen_run_idx]
205 | 
206 |         run_length = chosen_run.shape[0]
207 |         anchor_idx = self._rng.integers(
208 |             low=0,
209 |             high=run_length - self.min_distance,
210 |         )
211 |         end_idx = min(run_length, anchor_idx + self._max_rul)
212 |         query_idx = self._rng.integers(
213 |             low=anchor_idx + self.min_distance,
214 |             high=end_idx,
215 |         )
216 |         distance = query_idx - anchor_idx
217 | 
218 |         return chosen_run, anchor_idx, query_idx, distance, domain_label
219 | 
220 |     def _get_pair_idx_piecewise(self) -> Tuple[torch.Tensor, int, int, int, int]:
221 |         chosen_run_idx = self._rng.integers(0, len(self._features))
222 |         domain_label = self._run_domain_idx[chosen_run_idx]
223 |         chosen_run = self._features[chosen_run_idx]
224 | 
225 |         run_length = chosen_run.shape[0]
226 |         middle_idx = run_length // 2
227 |         anchor_idx = self._rng.integers(
228 |             low=0,
229 |             high=run_length - self.min_distance,
230 |         )
231 |         end_idx = (
232 |             middle_idx if anchor_idx < (middle_idx - self.min_distance) else run_length
233 |         )
234 |         query_idx = self._rng.integers(
235 |             low=anchor_idx + self.min_distance,
236 |             high=end_idx,
237 |         )
238 |         distance = query_idx - anchor_idx if anchor_idx > middle_idx else 0
239 | 
240 |         return chosen_run, anchor_idx, query_idx, distance, domain_label
241 | 
242 |     def _get_labeled_pair_idx(self) -> Tuple[torch.Tensor, int, int, int, int]:
243 |         chosen_run_idx = self._rng.integers(0, len(self._features))
244 |         domain_label = self._run_domain_idx[chosen_run_idx]
245 |         chosen_run = self._features[chosen_run_idx]
246 |         chosen_labels = self._labels[chosen_run_idx]
247 | 
248 |         run_length = chosen_run.shape[0]
249 |         anchor_idx = self._rng.integers(
250 |             low=0,
251 |             high=run_length - self.min_distance,
252 |         )
253 |         query_idx = self._rng.integers(
254 |             low=anchor_idx + self.min_distance,
255 |             high=run_length,
256 |         )
257 |         # RUL label difference is negative time step difference
258 |         distance = chosen_labels[anchor_idx] - chosen_labels[query_idx]
259 | 
260 |         return chosen_run, anchor_idx, query_idx, distance, domain_label
261 | 
262 |     def _build_pair(
263 |         self,
264 |         run: torch.Tensor,
265 |         anchor_idx: int,
266 |         query_idx: int,
267 |         distance: int,
268 |         domain_label: int,
269 |     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
270 |         anchors = run[anchor_idx]
271 |         queries = run[query_idx]
272 |         domain_label = torch.tensor(domain_label, dtype=torch.float)
273 |         distances = torch.tensor(distance, dtype=torch.float) / self._max_rul
274 |         distances = torch.clamp_max(distances, max=1)  # max distance is max_rul
275 | 
276 |         return anchors, queries, distances, domain_label
277 | 
278 | 
279 | class AdaptionDataset(Dataset):
280 |     def __init__(self, source, target, deterministic=False):
281 |         self.source = source
282 |         self.target = target
283 |         self.deterministic = deterministic
284 |         self._target_len = len(target)
285 | 
286 |         self._rng = np.random.default_rng(seed=42)
287 |         if self.deterministic:
288 |             self._get_target_idx = self._get_deterministic_target_idx
289 |             self._target_idx = [self._get_random_target_idx(_) for _ in range(len(self))]
290 |         else:
291 |             self._get_target_idx = self._get_random_target_idx
292 |             self._target_idx = None
293 | 
294 |     def _get_random_target_idx(self, _):
295 |         return self._rng.integers(0, self._target_len)
296 | 
297 |     def _get_deterministic_target_idx(self, idx):
298 |         return self._target_idx[idx]
299 | 
300 |     def __getitem__(self, idx):
301 |         target_idx = self._get_target_idx(idx)
302 |         source, source_label = self.source[idx]
303 |         target, _ = self.target[target_idx]
304 | 
305 |         return source, source_label, target
306 | 
307 |     def __len__(self):
308 |         return len(self.source)
309 | 


--------------------------------------------------------------------------------
/src/evaluation/add_rul_score.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | import re
  3 | 
  4 | import mlflow
  5 | import torch
  6 | import tqdm
  7 | 
  8 | import datasets
  9 | from lightning import metrics
 10 | from lightning.baseline import Baseline
 11 | 
 12 | LIST_PATTERN = re.compile(r"\d{1,3}")
 13 | 
 14 | 
 15 | def add_rul_score_all_fd(mlflow_uri):
 16 |     client = mlflow.tracking.MlflowClient(mlflow_uri)
 17 | 
 18 |     experiments = client.list_experiments()
 19 |     exp_pattern = re.compile("cmapss_.+_baseline$")
 20 |     experiments = [e for e in experiments if exp_pattern.match(e.name) is not None]
 21 |     print("Found experiments %d" % len(experiments))
 22 | 
 23 |     for e in experiments:
 24 |         runs = client.search_runs(e.experiment_id)
 25 | 
 26 |         errors = []
 27 |         for r in tqdm.tqdm(runs, desc=e.name):
 28 |             try:
 29 |                 _process_run(client, r)
 30 |             except Exception as e:
 31 |                 errors.append(e)
 32 | 
 33 |         print("Errors:")
 34 |         for e in errors:
 35 |             print(e)
 36 | 
 37 | 
 38 | def _process_run(client, r):
 39 |     # if run complete and not already processed
 40 |     if (
 41 |         "epoch" in r.data.metrics
 42 |         and r.data.metrics["epoch"] == 99.0
 43 |         and "val/rul_score" not in r.data.metrics
 44 |     ):
 45 |         dm = _get_datamodule(r)
 46 |         model_artifact = client.list_artifacts(r.info.run_id, "checkpoints")[0]
 47 |         model = _load_model(model_artifact, r)
 48 |         score = _get_score(model, dm.val_dataloader())
 49 |         _log_metric(client, r, score, "val/rul_score")
 50 |         for fd, test_loader in enumerate(dm.test_dataloader(), start=1):
 51 |             score = _get_score(model, test_loader)
 52 |             _log_metric(client, r, score, f"test/rul_score_fd{fd}")
 53 | 
 54 | 
 55 | def _get_datamodule(r):
 56 |     data = datasets.BaselineDataModule(
 57 |         fd_source=int(r.data.params["fd_source"]),
 58 |         batch_size=512,
 59 |     )
 60 |     data.prepare_data()
 61 |     data.setup()
 62 | 
 63 |     return data
 64 | 
 65 | 
 66 | def _load_model(model_artifact, r):
 67 |     artifact_path = os.path.join(
 68 |         r.info.artifact_uri.replace("file://", ""), model_artifact.path
 69 |     )
 70 |     model = Baseline.load_from_checkpoint(artifact_path, map_location="cpu")
 71 |     if "pretrained_checkpoint" in r.data.params:
 72 |         model.encoder.norm_outputs = True
 73 |     model.eval()
 74 | 
 75 |     return model
 76 | 
 77 | 
 78 | @torch.no_grad()
 79 | def _get_score(model, dataloader):
 80 |     model = model.to("cuda:0")
 81 |     rul_score = metrics.RULScore()
 82 |     rul_score_metric = metrics.SimpleMetric(reduction="sum")
 83 |     for batch in dataloader:
 84 |         features, labels = [b.to("cuda:0") for b in batch]
 85 |         batch_size = len(features)
 86 |         score = rul_score(model(features), labels)
 87 |         rul_score_metric.update(score, batch_size)
 88 | 
 89 |     return rul_score_metric.compute().item()
 90 | 
 91 | 
 92 | def _log_metric(client, r, score, metric):
 93 |     if metric not in r.data.metrics:
 94 |         client.log_metric(r.info.run_id, metric, score)
 95 | 
 96 | 
 97 | if __name__ == "__main__":
 98 |     import argparse
 99 | 
100 |     parser = argparse.ArgumentParser(
101 |         description="Add RUL score val and test metric for all FDs of CMAPSS"
102 |     )
103 |     parser.add_argument("mlflow_uri", help="tracking URI for mlflow")
104 |     opt = parser.parse_args()
105 | 
106 |     add_rul_score_all_fd(opt.mlflow_uri)
107 | 


--------------------------------------------------------------------------------
/src/evaluation/export_results_semi_supervised.py:
--------------------------------------------------------------------------------
  1 | """Exports all RUL experiments to a data frame."""
  2 | 
  3 | import argparse
  4 | import re
  5 | 
  6 | import mlflow
  7 | import numpy as np
  8 | import pandas as pd
  9 | 
 10 | LIST_PATTERN = re.compile(r".?\d{1,3}")
 11 | VERSION_PATTERN = re.compile(r".*?(@(?P<percent_broken>.*))?@(?P<percent_fail_runs>.*)")
 12 | 
 13 | 
 14 | def export_all(mlflow_uri, tags):
 15 |     """Export all transfer and baseline experiments to CSV."""
 16 |     client = mlflow.tracking.MlflowClient(mlflow_uri)
 17 | 
 18 |     print("Export baseline experiments...")
 19 |     df = export_matching_experiments(
 20 |         client, "cmapss_.{3,5}_baseline$", tags, _runs_of_semi_supervised
 21 |     )
 22 |     if df is not None:
 23 |         df.to_csv("semi_supervised.csv")
 24 | 
 25 | 
 26 | def export_matching_experiments(client, exp_name_regex, tags, func):
 27 |     regex = re.compile(exp_name_regex)
 28 |     experiments = client.list_experiments()
 29 |     filtered_experiments = [e for e in experiments if regex.match(e.name) is not None]
 30 |     print("Found %d experiments..." % len(filtered_experiments))
 31 |     df = []
 32 |     for e in filtered_experiments:
 33 |         print('Evaluate experiment "%s" with id %s' % (e.name, e.experiment_id))
 34 |         df.append(func(client, e, tags))
 35 |     if df:
 36 |         df = pd.concat(df)
 37 |     else:
 38 |         df = None
 39 | 
 40 |     return df
 41 | 
 42 | 
 43 | def _runs_of_semi_supervised(client, e, tags):
 44 |     replications = _get_replications(client, e)
 45 |     replications = _filter_complete_with_tag(replications, tags)
 46 |     df = pd.DataFrame(
 47 |         np.zeros((len(replications), 7)),
 48 |         columns=[
 49 |             "source",
 50 |             "num_labeled",
 51 |             "percent_broken",
 52 |             "pretrained",
 53 |             "test",
 54 |             "val",
 55 |             "score",
 56 |         ],
 57 |         index=[e.name] * len(replications),
 58 |     )
 59 |     for i, run in enumerate(replications):
 60 |         statics = [
 61 |             run.data.params["fd_source"],
 62 |             _get_num_labeled(run.data.params["percent_fail_runs"]),
 63 |             _get_percent_broken(run.data.tags["version"]),
 64 |             _get_pretraining_mode(run),
 65 |         ]
 66 |         test_rmse = _get_test_value(f"regression_loss_fd{statics[0]}", run)
 67 |         test_score = _get_test_value(f"rul_score_fd{statics[0]}", run)
 68 |         val_rmse = _get_val_value(client, "regression_loss", "regression_loss", run)
 69 |         df.iloc[i] = statics + [test_rmse, val_rmse, test_score]
 70 |     print("Return %d runs..." % len(df))
 71 | 
 72 |     return df
 73 | 
 74 | 
 75 | def _get_replications(client, experiment):
 76 |     runs = client.search_runs(experiment_ids=experiment.experiment_id, max_results=9999)
 77 |     print("Found %d top-level runs..." % len(runs))
 78 | 
 79 |     return runs
 80 | 
 81 | 
 82 | def _filter_complete_with_tag(replications, tags):
 83 |     versions = [
 84 |         r.data.tags["version"]
 85 |         for r in replications
 86 |         if any(r.data.tags["version"].startswith(t) for t in tags)
 87 |     ]
 88 |     versions, version_counts = np.unique(versions, return_counts=True)
 89 |     versions = versions[version_counts == 10]
 90 |     replications = [r for r in replications if r.data.tags["version"] in versions]
 91 | 
 92 |     return replications
 93 | 
 94 | 
 95 | def _get_num_labeled(percent_labeled):
 96 |     return len(LIST_PATTERN.findall(percent_labeled))
 97 | 
 98 | 
 99 | def _get_pretraining_mode(run):
100 |     version = run.data.tags["version"]
101 |     if "baseline" in version:
102 |         return "no pretraining"
103 |     elif "autoencoder" in version:
104 |         return "autoencoder"
105 |     elif "rbm" in version:
106 |         return "rbm"
107 |     else:
108 |         return "metric"
109 | 
110 | 
111 | def _get_percent_broken(version):
112 |     match = VERSION_PATTERN.match(version)
113 |     if match is not None:
114 |         percent_broken = match.group("percent_broken") or "0.0"
115 |         percent_broken = float(percent_broken)
116 |     else:
117 |         raise ValueError(f"Could not parse version tag '{version}'")
118 | 
119 |     return percent_broken
120 | 
121 | 
122 | def _get_test_value(metric, run):
123 |     """Return test value of selected metric."""
124 |     test_metric = f"test/{metric}"
125 |     if test_metric in run.data.metrics:
126 |         best_value = run.data.metrics[test_metric]
127 |     else:
128 |         best_value = np.nan
129 | 
130 |     return best_value
131 | 
132 | 
133 | def _get_val_value(client, metric, indicator_metric, run):
134 |     """Return test value of selected metric."""
135 |     val_metric = f"val/{metric}"
136 |     val_indicator = f"val/{indicator_metric}"
137 |     if val_metric in run.data.metrics and val_indicator in run.data.metrics:
138 |         indicator_values = client.get_metric_history(run.info.run_id, val_indicator)
139 |         best_indicator_idx = np.argmin([m.value for m in indicator_values])
140 |         metric_value = client.get_metric_history(run.info.run_id, val_metric)
141 |         best_value = metric_value[best_indicator_idx].value
142 |     else:
143 |         best_value = np.nan
144 | 
145 |     return best_value
146 | 
147 | 
148 | if __name__ == "__main__":
149 |     parser = argparse.ArgumentParser(
150 |         description="Exports semi_supervised experiments to CSV"
151 |     )
152 |     parser.add_argument("mlflow_uri", help="URI for MLFlow Server.")
153 |     parser.add_argument(
154 |         "--tags", "-t", nargs="+", help="MLFlow tag prefix of runs to export"
155 |     )
156 |     opt = parser.parse_args()
157 | 
158 |     export_all(opt.mlflow_uri, opt.tags)
159 | 


--------------------------------------------------------------------------------
/src/evaluation/plot_semi_supervised.py:
--------------------------------------------------------------------------------
  1 | import re
  2 | 
  3 | import numpy as np
  4 | import pandas as pd
  5 | import plotnine
  6 | from matplotlib import cm, colors
  7 | from plotnine import aes, ggplot
  8 | from tabulate import tabulate
  9 | 
 10 | 
 11 | def plot(file_path):
 12 |     df = pd.read_csv(file_path, index_col=0)
 13 |     df = _preprocess_data(df)
 14 | 
 15 |     percent_broken = [40, 60, 70, 80, 90]
 16 |     percent_labeled = [2, 10, 20, 40, 100]
 17 | 
 18 |     plotnine.options.figure_size = (15, 8)
 19 |     for broken in percent_broken:
 20 |         plot_df = df[(df["percent_broken"] == broken) | (df["percent_broken"] == 0.0)]
 21 |         _plot_val_test(plot_df, broken, file_path)
 22 |     for labeled in percent_labeled:
 23 |         plot_df = df[(df["num_labeled"] == labeled) & (df["percent_broken"] > 0.0)]
 24 |         gg = _vs_broken(plot_df, "test")
 25 |         gg.save(file_path.replace(".csv", f"_test_labeled@{labeled}.pdf"))
 26 |         gg = _vs_broken(plot_df, "score", log_scale=True)
 27 |         gg.save(file_path.replace(".csv", f"_score_labeled@{labeled}.pdf"))
 28 |     _generate_table(df, "test")
 29 |     _generate_table(df, "score")
 30 | 
 31 | 
 32 | def _preprocess_data(df):
 33 |     method2displayname = {
 34 |         "no pretraining": "None",
 35 |         "metric": "Self-Supervised",
 36 |         "autoencoder": "Autoencoder",
 37 |         "rbm": "RBM",
 38 |     }
 39 |     df["Pre-Training"] = df["pretrained"].map(method2displayname)
 40 |     num2percent_labeled = {
 41 |         1.0: 2,
 42 |         4.0: 2,
 43 |         3.0: 2,
 44 |         8.0: 10,
 45 |         20.0: 10,
 46 |         19.0: 10,
 47 |         16.0: 20,
 48 |         41.0: 20,
 49 |         39.0: 20,
 50 |         32.0: 40,
 51 |         83.0: 40,
 52 |         79.0: 40,
 53 |         80.0: 100,
 54 |         208.0: 100,
 55 |         199.0: 100,
 56 |     }
 57 |     df["num_labeled"] = df["num_labeled"].map(num2percent_labeled)
 58 |     num2subset = {1: "FD001", 2: "FD002", 3: "FD003", 4: "FD004"}
 59 |     df["source"] = df["source"].map(num2subset)
 60 |     df["percent_broken"] = (df["percent_broken"] * 100).astype(int)
 61 | 
 62 |     return df
 63 | 
 64 | 
 65 | def _plot_val_test(plot_df, broken, file_path):
 66 |     plot_df.loc[:, "Pre-Training"] = pd.Categorical(
 67 |         plot_df["Pre-Training"],
 68 |         ["None", "Self-Supervised", "Autoencoder", "RBM"],
 69 |         ordered=True,
 70 |     )
 71 |     gg = _vs_labeled(plot_df, "val")
 72 |     gg.save(file_path.replace(".csv", f"_val@{broken:.2f}.pdf"))
 73 |     gg = _vs_labeled(plot_df, "test")
 74 |     gg.save(file_path.replace(".csv", f"_test@{broken:.2f}.pdf"))
 75 |     gg = _vs_labeled(plot_df, "score", log_scale=True)
 76 |     gg.save(file_path.replace(".csv", f"_score@{broken:.2f}.pdf"))
 77 | 
 78 | 
 79 | def _vs_labeled(df, column, log_scale=False):
 80 |     df = df.sort_values("Pre-Training")
 81 |     gg = (
 82 |         ggplot(df, aes(x="num_labeled", y=column))
 83 |         + plotnine.stat_boxplot(
 84 |             aes(
 85 |                 y=column,
 86 |                 x="factor(num_labeled)",
 87 |                 color="Pre-Training",
 88 |             )
 89 |         )
 90 |         + plotnine.facet_wrap("source", nrow=2, ncol=2, scales="free_x")
 91 |         + plotnine.scale_color_manual(
 92 |             [colors.rgb2hex(c) for c in cm.get_cmap("tab10").colors]
 93 |         )
 94 |         + plotnine.xlab("Percent Labeled")
 95 |         + plotnine.theme_classic(base_size=20)
 96 |         + plotnine.theme(subplots_adjust={"hspace": 0.25})
 97 |     )
 98 |     if column == "score":
 99 |         gg += plotnine.ylab("RUL Score")
100 |     else:
101 |         gg += plotnine.ylab("RMSE")
102 |     if log_scale:
103 |         gg += plotnine.scale_y_log10()
104 | 
105 |     return gg
106 | 
107 | 
108 | def _vs_broken(df, column, log_scale=False):
109 |     df.loc[:, "Pre-Training"] = pd.Categorical(
110 |         df["Pre-Training"],
111 |         ["Self-Supervised", "Autoencoder", "RBM"],
112 |         ordered=True,
113 |     )
114 |     df = df.sort_values("Pre-Training")
115 |     gg = (
116 |         ggplot(df, aes(x="percent_broken", y=column))
117 |         + plotnine.stat_boxplot(
118 |             aes(
119 |                 y=column,
120 |                 x="factor(percent_broken)",
121 |                 color="Pre-Training",
122 |             )
123 |         )
124 |         + plotnine.facet_wrap("source", nrow=2, ncol=2)
125 |         + plotnine.scale_color_manual(
126 |             [colors.rgb2hex(c) for c in cm.get_cmap("tab10").colors[1:]]
127 |         )
128 |         + plotnine.xlab("Grade of Degradation")
129 |         + plotnine.theme_classic(base_size=20)
130 |         + plotnine.theme(subplots_adjust={"hspace": 0.25})
131 |     )
132 |     if column == "score":
133 |         gg += plotnine.ylab("RUL Score")
134 |     else:
135 |         gg += plotnine.ylab("RMSE")
136 |     if log_scale:
137 |         gg += plotnine.scale_y_log10()
138 | 
139 |     return gg
140 | 
141 | 
142 | POWER_PATTERN = re.compile(r"(?P<base>[\d.]+)e(?P<exponent>[+-]\d+)")
143 | 
144 | 
145 | def _replace_power_fmt(string):
146 |     for match in POWER_PATTERN.finditer(string):
147 |         base = match.group("base")
148 |         exponent = int(match.group("exponent"))
149 |         string = string.replace(match.group(), fr"{base}e{exponent}")
150 | 
151 |     return string
152 | 
153 | 
154 | def _generate_table(df: pd.DataFrame, column):
155 |     df = df.filter(
156 |         items=["source", "Pre-Training", "percent_broken", "num_labeled", column],
157 |         axis=1,
158 |     )
159 |     df = df.groupby(["source", "Pre-Training", "percent_broken", "num_labeled"])
160 |     df = df.agg(
161 |         mean=pd.NamedAgg(column, "mean"), std=pd.NamedAgg(column, "std")
162 |     ).reset_index()
163 |     if column == "score":
164 |         df[column] = df["mean"].combine(
165 |             df["std"], func=lambda mean, std: f"${mean:.2e} \pm {std:.2e}$"
166 |         )
167 |         df[column] = df[column].apply(_replace_power_fmt)
168 |     else:
169 |         df[column] = df["mean"].combine(
170 |             df["std"], func=lambda mean, std: f"${mean:.2f} \pm {std:.2f}$"
171 |         )
172 |     argmins = _get_argmins(df)
173 |     df = df.pivot(
174 |         index=["Pre-Training", "percent_broken"],
175 |         columns=["source", "num_labeled"],
176 |         values=[column],
177 |     )
178 |     df = (
179 |         df.reset_index(level="Pre-Training")
180 |         .sort_index(level="percent_broken")
181 |         .droplevel(0, axis=1)
182 |     )
183 |     df = _mark_best(df, argmins)
184 |     df = _rename_methods(df)
185 |     df = _make_multirows(df)
186 | 
187 |     for fd in ["FD001", "FD002", "FD003", "FD004"]:
188 |         print(tabulate(df.loc[:, ["", fd]], headers="keys", tablefmt="latex_raw"))
189 | 
190 | 
191 | def _get_argmins(df):
192 |     min_df = df.pivot(
193 |         index=["Pre-Training", "percent_broken"],
194 |         columns=["source", "num_labeled"],
195 |         values=["mean"],
196 |     )
197 |     worse_than_baseline = (min_df - min_df.loc[("None", 0), :]) > 0
198 |     worse_than_baseline = worse_than_baseline.groupby(["percent_broken"]).agg(all)
199 |     argmins = min_df.reset_index().groupby(["percent_broken"]).agg(np.argmin)
200 |     argmins = argmins.drop(("Pre-Training", "", ""), axis=1)
201 |     argmins[worse_than_baseline] = -1
202 |     argmins = argmins.droplevel(0, axis=1)
203 | 
204 |     return argmins
205 | 
206 | 
207 | def _mark_best(df, argmins):
208 |     row_iter = argmins.iterrows()
209 |     next(row_iter)  # skip baseline row
210 |     for percent_broken, row in row_iter:
211 |         for col_idx in row.index:
212 |             min_idx = int(argmins.loc[percent_broken, col_idx])  #
213 |             if min_idx > -1:
214 |                 stripped_value = df.loc[percent_broken, col_idx].iloc[min_idx].strip("$")
215 |                 df.loc[percent_broken, col_idx].iloc[
216 |                     min_idx
217 |                 ] = f"$\mathbf{{{stripped_value}}}$"
218 | 
219 |     return df
220 | 
221 | 
222 | def _rename_methods(df):
223 |     method_dict = {
224 |         "Autoencoder": "AE",
225 |         "Self-Supervised": "Ours",
226 |         "RBM": "RBM",
227 |         "None": "None",
228 |     }
229 |     df.loc[:, ("", "")] = df.loc[:, ("", "")].map(method_dict)
230 | 
231 |     return df
232 | 
233 | 
234 | def _make_multirows(df):
235 |     new_index = []
236 |     for i, value in enumerate(df.index):
237 |         if value == 0:
238 |             new_index.append("")
239 |         elif i % 3 == 1:
240 |             new_index.append(f"\multirow{{3}}{{*}}{{{value}\%}}")
241 |         else:
242 |             new_index.append("")
243 |     df.index = new_index
244 | 
245 |     return df
246 | 
247 | 
248 | if __name__ == "__main__":
249 |     import argparse
250 | 
251 |     parser = argparse.ArgumentParser(
252 |         description="Plot results of semi-supervised experiments"
253 |     )
254 |     parser.add_argument("file_path", help="path to result CSV file")
255 |     opt = parser.parse_args()
256 | 
257 |     plot(opt.file_path)
258 | 


--------------------------------------------------------------------------------
/src/hyperopt/hyperopt_semi_supervised.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | from functools import partial
  3 | 
  4 | import numpy as np
  5 | import pytorch_lightning as pl
  6 | import pytorch_lightning.loggers as pl_loggers
  7 | from ray import tune
  8 | 
  9 | import building
 10 | import datasets
 11 | from lightning import loggers
 12 | 
 13 | 
 14 | def tune_pretraining(config, arch_config, source, percent_broken, encoder, mode):
 15 |     best_scores = []
 16 |     for i in range(5):
 17 |         logger = pl_loggers.TensorBoardLogger(
 18 |             _get_hyperopt_logdir(),
 19 |             loggers.semi_supervised_hyperopt_name(source, percent_broken),
 20 |         )
 21 |         checkpoint_callback = pl.callbacks.ModelCheckpoint(monitor="val/checkpoint_score")
 22 |         trainer = building.build_trainer(
 23 |             logger,
 24 |             checkpoint_callback,
 25 |             max_epochs=100,
 26 |             val_interval=1.0,
 27 |             gpu=1,
 28 |             seed=None,
 29 |             check_sanity=False,
 30 |         )
 31 | 
 32 |         data = datasets.PretrainingBaselineDataModule(
 33 |             fd_source=source,
 34 |             num_samples=50000,
 35 |             batch_size=config["batch_size"],
 36 |             percent_broken=percent_broken,
 37 |             percent_fail_runs=0.0,
 38 |             min_distance=config["min_distance"],
 39 |             truncate_val=True,
 40 |         )
 41 |         if mode == "metric":
 42 |             model = building.build_pretraining_from_config(
 43 |                 arch_config,
 44 |                 config,
 45 |                 data.window_size,
 46 |                 encoder=encoder,
 47 |                 record_embeddings=False,
 48 |                 use_adaption=False,
 49 |             )
 50 |         else:
 51 |             model = building.build_autoencoder_from_config(
 52 |                 arch_config,
 53 |                 config,
 54 |                 data.window_size,
 55 |                 encoder=encoder,
 56 |                 record_embeddings=False,
 57 |                 use_adaption=False,
 58 |             )
 59 |         building.add_hparams(model, data, 42)
 60 | 
 61 |         trainer.fit(model, datamodule=data)
 62 |         best_scores.append(trainer.checkpoint_callback.best_model_score.item())
 63 |         if np.std(best_scores) > 0.005 or np.max(best_scores) > 0.04:
 64 |             tune.report(checkpoint_score=np.max(best_scores), replication=i)
 65 |             return
 66 | 
 67 |     tune.report(checkpoint_score=np.mean(best_scores), replication=5)
 68 | 
 69 | 
 70 | def _get_hyperopt_logdir():
 71 |     script_path = os.path.dirname(__file__)
 72 |     log_dir = os.path.normpath(os.path.join(script_path, "..", "..", "hyperopt"))
 73 | 
 74 |     return log_dir
 75 | 
 76 | 
 77 | def optimize_pretraining(source, percent_broken, arch_config, encoder, mode, num_trials):
 78 |     config = {
 79 |         "domain_tradeoff": tune.choice([0.0]),
 80 |         "dropout": tune.quniform(0.0, 0.5, 0.1),
 81 |         "lr": tune.qloguniform(1e-4, 1e-1, 5e-5),
 82 |         "batch_size": tune.choice([64, 128, 256, 512]),
 83 |         "min_distance": tune.choice([1, 10, 15, 30]),
 84 |     }
 85 | 
 86 |     scheduler = tune.schedulers.FIFOScheduler()
 87 |     reporter = tune.CLIReporter(
 88 |         parameter_columns=list(config.keys()),
 89 |         metric_columns=["checkpoint_score", "replication"],
 90 |     )
 91 | 
 92 |     tune_func = partial(
 93 |         tune_pretraining,
 94 |         arch_config=arch_config,
 95 |         source=source,
 96 |         percent_broken=percent_broken,
 97 |         encoder=encoder,
 98 |         mode=mode,
 99 |     )
100 |     analysis = tune.run(
101 |         tune_func,
102 |         resources_per_trial={"cpu": 3, "gpu": 0.5},
103 |         metric="checkpoint_score",
104 |         mode="min",
105 |         config=config,
106 |         num_samples=num_trials,
107 |         scheduler=scheduler,
108 |         progress_reporter=reporter,
109 |         name="tune_stability_asha",
110 |         fail_fast=True,
111 |     )
112 | 
113 |     print("Best hyperparameters found were: ", analysis.best_config)
114 | 
115 |     return analysis.best_config
116 | 
117 | 
118 | if __name__ == "__main__":
119 |     import argparse
120 | 
121 |     parser = argparse.ArgumentParser(
122 |         description="Hyperparameter optimization for pretraining"
123 |     )
124 |     parser.add_argument("--source", type=int, required=True, help="source FD number")
125 |     parser.add_argument(
126 |         "--percent_broken", type=float, required=True, help="degradation in [0, 1]"
127 |     )
128 |     parser.add_argument("arch_config_path", help="path to architecture config JSON")
129 |     parser.add_argument(
130 |         "--encoder",
131 |         default="cnn",
132 |         choices=["cnn", "lstm"],
133 |         help="encoder type",
134 |     )
135 |     parser.add_argument(
136 |         "--mode",
137 |         choices=["metric", "autoencoder"],
138 |         default="metric",
139 |         help="metric or autoencoder pretraining mode",
140 |     )
141 |     parser.add_argument(
142 |         "--num_trials", type=int, required=True, help="number of hyperopt trials"
143 |     )
144 |     opt = parser.parse_args()
145 | 
146 |     _arch_config = building.load_config(opt.arch_config_path)
147 |     optimize_pretraining(
148 |         opt.source,
149 |         opt.percent_broken,
150 |         _arch_config,
151 |         opt.encoder,
152 |         opt.mode,
153 |         opt.num_trials,
154 |     )
155 | 


--------------------------------------------------------------------------------
/src/hyperopt/hyperopt_supervised.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | from functools import partial
  3 | 
  4 | import pytorch_lightning as pl
  5 | import pytorch_lightning.loggers as pl_loggers
  6 | from ray import tune
  7 | from ray.tune.integration.pytorch_lightning import TuneReportCallback
  8 | 
  9 | import building
 10 | import datasets
 11 | from lightning import loggers
 12 | 
 13 | 
 14 | def tune_supervised(config, source, arch_config, encoder):
 15 |     arch_config.update(config)
 16 | 
 17 |     logger = pl_loggers.TensorBoardLogger(
 18 |         _get_hyperopt_logdir(),
 19 |         loggers.baseline_experiment_name(source),
 20 |     )
 21 |     checkpoint_callback = pl.callbacks.ModelCheckpoint(monitor="val/regression_loss")
 22 |     tune_callback = TuneReportCallback(
 23 |         {
 24 |             "reg_loss": "val/regression_loss",
 25 |         },
 26 |         on="validation_end",
 27 |     )
 28 |     trainer = building.build_trainer(
 29 |         logger,
 30 |         checkpoint_callback,
 31 |         max_epochs=100,
 32 |         val_interval=1.0,
 33 |         gpu=1,
 34 |         seed=None,
 35 |         callbacks=[tune_callback],
 36 |         check_sanity=False,
 37 |     )
 38 | 
 39 |     data = datasets.BaselineDataModule(
 40 |         fd_source=source, batch_size=arch_config["batch_size"]
 41 |     )
 42 |     model = building.build_baseline_from_config(
 43 |         arch_config, data.window_size, encoder, None
 44 |     )
 45 |     building.add_hparams(model, data, None)
 46 | 
 47 |     trainer.fit(model, datamodule=data)
 48 | 
 49 | 
 50 | def _get_hyperopt_logdir():
 51 |     script_path = os.path.dirname(__file__)
 52 |     log_dir = os.path.normpath(os.path.join(script_path, "..", "..", "hyperopt"))
 53 | 
 54 |     return log_dir
 55 | 
 56 | 
 57 | def optimize_supervised(source, arch_config, encoder, num_trials):
 58 |     config = {
 59 |         "dropout": tune.quniform(0.0, 0.5, 0.1),
 60 |         "lr": tune.qloguniform(1e-4, 1e-1, 5e-5),
 61 |         "batch_size": tune.choice([64, 128, 256, 512]),
 62 |     }
 63 | 
 64 |     scheduler = tune.schedulers.ASHAScheduler(
 65 |         max_t=100, grace_period=10, reduction_factor=2
 66 |     )
 67 |     reporter = tune.CLIReporter(
 68 |         parameter_columns=list(config.keys()),
 69 |         metric_columns=["reg_loss"],
 70 |     )
 71 | 
 72 |     tune_func = partial(
 73 |         tune_supervised,
 74 |         source=source,
 75 |         arch_config=arch_config,
 76 |         encoder=encoder,
 77 |     )
 78 |     analysis = tune.run(
 79 |         tune_func,
 80 |         resources_per_trial={"cpu": 3, "gpu": 0.5},
 81 |         metric="reg_loss",
 82 |         mode="min",
 83 |         config=config,
 84 |         num_samples=num_trials,
 85 |         scheduler=scheduler,
 86 |         progress_reporter=reporter,
 87 |         name="tune_supervised_asha",
 88 |         fail_fast=True,
 89 |     )
 90 | 
 91 |     print("Best hyperparameters found were: ", analysis.best_config)
 92 | 
 93 |     return analysis.best_config
 94 | 
 95 | 
 96 | if __name__ == "__main__":
 97 |     import argparse
 98 | 
 99 |     parser = argparse.ArgumentParser(
100 |         description="Hyperparameter optimization for supervised training"
101 |     )
102 |     parser.add_argument("--source", type=int, required=True, help="source FD number")
103 |     parser.add_argument(
104 |         "--arch_config", required=True, help="path to architecture base config"
105 |     )
106 |     parser.add_argument(
107 |         "--encoder",
108 |         default="cnn",
109 |         choices=["cnn", "lstm"],
110 |         help="encoder type",
111 |     )
112 |     parser.add_argument(
113 |         "--num_trials", type=int, required=True, help="number of hyperopt trials"
114 |     )
115 |     opt = parser.parse_args()
116 | 
117 |     _arch_config = building.load_config(opt.arch_config)
118 |     optimize_supervised(opt.source, _arch_config, opt.encoder, opt.num_trials)
119 | 


--------------------------------------------------------------------------------
/src/lightning/autoencoder.py:
--------------------------------------------------------------------------------
  1 | import pytorch_lightning as pl
  2 | import torch
  3 | import torch.nn as nn
  4 | 
  5 | from lightning import metrics
  6 | from lightning.mixins import DataHparamsMixin
  7 | from models import networks
  8 | 
  9 | 
 10 | class AutoencoderPretraining(pl.LightningModule, DataHparamsMixin):
 11 |     def __init__(
 12 |         self,
 13 |         in_channels,
 14 |         seq_len,
 15 |         num_layers,
 16 |         kernel_size,
 17 |         base_filters,
 18 |         latent_dim,
 19 |         dropout,
 20 |         domain_tradeoff,
 21 |         domain_disc_dim,
 22 |         num_disc_layers,
 23 |         lr,
 24 |         weight_decay,
 25 |         record_embeddings=False,
 26 |         encoder="cnn",
 27 |     ):
 28 |         super().__init__()
 29 | 
 30 |         self.in_channels = in_channels
 31 |         self.seq_len = seq_len
 32 |         self.num_layers = num_layers
 33 |         self.kernel_size = kernel_size
 34 |         self.base_filters = base_filters
 35 |         self.latent_dim = latent_dim
 36 |         self.dropout = dropout
 37 |         self.domain_tradeoff = domain_tradeoff
 38 |         self.domain_disc_dim = domain_disc_dim
 39 |         self.num_disc_layers = num_disc_layers
 40 |         self.lr = lr
 41 |         self.weight_decay = weight_decay
 42 |         self.record_embeddings = record_embeddings
 43 | 
 44 |         self.encoder = networks.Encoder(
 45 |             self.in_channels,
 46 |             self.base_filters,
 47 |             self.kernel_size,
 48 |             self.num_layers,
 49 |             self.latent_dim,
 50 |             self.seq_len,
 51 |             self.dropout,
 52 |             norm_outputs=True,
 53 |         )
 54 |         self.decoder = networks.Decoder(
 55 |             self.in_channels,
 56 |             self.base_filters,
 57 |             self.kernel_size,
 58 |             self.num_layers,
 59 |             self.latent_dim,
 60 |             self.seq_len,
 61 |             self.dropout,
 62 |         )
 63 |         self.domain_disc = None
 64 | 
 65 |         self.criterion_regression = nn.MSELoss()
 66 |         self.criterion_domain = nn.BCEWithLogitsLoss()
 67 | 
 68 |         self.embedding_metric = metrics.EmbeddingViz(40000, self.latent_dim)
 69 |         self.regression_metric = metrics.SimpleMetric(40000)
 70 |         self.domain_metric = metrics.SimpleMetric(40000)
 71 | 
 72 |         self.save_hyperparameters()
 73 |         self.hparams["mode"] = "autoencoder"
 74 | 
 75 |     def configure_optimizers(self):
 76 |         return torch.optim.Adam(
 77 |             self.parameters(), lr=self.lr, weight_decay=self.weight_decay
 78 |         )
 79 | 
 80 |     def forward(self, inputs):
 81 |         latent_code = self.encoder(inputs)
 82 |         outputs = self.decoder(latent_code)
 83 | 
 84 |         return outputs
 85 | 
 86 |     def training_step(self, batch, batch_idx):
 87 |         regression_loss, domain_loss = self._get_losses(batch)
 88 |         loss = regression_loss + self.domain_tradeoff * domain_loss
 89 |         self.log("train/loss", loss)
 90 |         self.log("train/regression_loss", regression_loss)
 91 |         self.log("train/domain_loss", domain_loss)
 92 | 
 93 |         return loss
 94 | 
 95 |     def on_validation_epoch_start(self):
 96 |         self._reset_all_metrics()
 97 | 
 98 |     def _reset_all_metrics(self):
 99 |         self.embedding_metric.reset()
100 |         self.regression_metric.reset()
101 |         self.domain_metric.reset()
102 | 
103 |     def validation_step(self, batch, batch_idx, dataloader_idx):
104 |         if dataloader_idx == 0:
105 |             regression_loss, domain_loss = self._get_losses(batch)
106 |             batch_size = batch[0].shape[0]
107 |             self.regression_metric.update(regression_loss, batch_size)
108 |             self.domain_metric.update(domain_loss, batch_size)
109 |         elif self.record_embeddings:
110 |             self._record_embeddings(batch, dataloader_idx)
111 | 
112 |     def _record_embeddings(self, batch, dataloader_idx):
113 |         features, labels = batch
114 |         embedding = self.encoder(features)
115 |         domain_labels = torch.full_like(
116 |             labels, fill_value=(2 - dataloader_idx), dtype=torch.int
117 |         )
118 |         self.embedding_metric.update(embedding, domain_labels, labels)
119 | 
120 |     def validation_epoch_end(self, validation_step_outputs):
121 |         if self.record_embeddings:
122 |             embedding_fig = self.embedding_metric.compute()
123 |             self.logger.log_figure("val/embeddings", embedding_fig, self.global_step)
124 | 
125 |         regression_loss = self.regression_metric.compute()
126 |         domain_loss = self.domain_metric.compute()
127 | 
128 |         self.log("val/regression_loss", regression_loss)
129 |         self.log("val/domain_loss", domain_loss)
130 |         self.log("val/checkpoint_score", regression_loss - 0.1 * domain_loss)
131 | 
132 |     def _get_losses(self, batch):
133 |         anchors, queries, true_distances, domain_labels = batch
134 |         combined = torch.cat([anchors, queries])
135 |         embeddings = self.encoder(combined)
136 |         outputs = self.decoder(embeddings)
137 |         regression_loss = self.criterion_regression(outputs, combined)
138 | 
139 |         if self.domain_tradeoff > 0:
140 |             batch_size = anchors.shape[0]
141 |             domain_pred = self.domain_disc(embeddings[:batch_size])
142 |             domain_loss = self.criterion_domain(domain_pred, domain_labels)
143 |         else:
144 |             domain_loss = 0
145 | 
146 |         return regression_loss, domain_loss
147 | 


--------------------------------------------------------------------------------
/src/lightning/baseline.py:
--------------------------------------------------------------------------------
  1 | import pytorch_lightning as pl
  2 | import torch
  3 | 
  4 | from lightning import metrics
  5 | from lightning.mixins import LoadEncoderMixin
  6 | from models import networks
  7 | 
  8 | 
  9 | class Baseline(pl.LightningModule, LoadEncoderMixin):
 10 |     def __init__(
 11 |         self,
 12 |         in_channels,
 13 |         seq_len,
 14 |         num_layers,
 15 |         kernel_size,
 16 |         base_filters,
 17 |         latent_dim,
 18 |         dropout=0.0,
 19 |         optim_type="adam",
 20 |         lr=0.001,
 21 |         record_embeddings=False,
 22 |         encoder="cnn",
 23 |     ):
 24 |         super().__init__()
 25 | 
 26 |         self.in_channels = in_channels
 27 |         self.seq_len = seq_len
 28 |         self.num_layers = num_layers
 29 |         self.kernel_size = kernel_size
 30 |         self.base_filters = base_filters
 31 |         self.latent_dim = latent_dim
 32 |         self.dropout = dropout
 33 |         self.optim_type = optim_type
 34 |         self.lr = lr
 35 |         self.record_embeddings = record_embeddings
 36 |         self.encoder_type = encoder
 37 | 
 38 |         self.encoder = self._get_encoder()
 39 |         self.regressor = networks.Regressor(latent_dim)
 40 | 
 41 |         self.criterion_regression = metrics.RMSELoss(num_elements=0)
 42 | 
 43 |         self.embedding_metric = metrics.EmbeddingViz(
 44 |             20000, self.latent_dim, combined=False
 45 |         )
 46 |         self.regression_metrics = {i: metrics.RMSELoss() for i in range(1, 5)}
 47 | 
 48 |         self.save_hyperparameters()
 49 | 
 50 |     def _get_encoder(self):
 51 |         if "cnn" in self.encoder_type:
 52 |             encoder = networks.Encoder(
 53 |                 self.in_channels,
 54 |                 self.base_filters,
 55 |                 self.kernel_size,
 56 |                 self.num_layers,
 57 |                 self.latent_dim,
 58 |                 self.seq_len,
 59 |                 dropout=self.dropout,
 60 |                 norm_outputs=False,
 61 |             )
 62 |         else:
 63 |             raise ValueError(
 64 |                 f"Unknown encoder type {self.encoder_type}. Must contain 'cnn'."
 65 |             )
 66 | 
 67 |         return encoder
 68 | 
 69 |     def add_data_hparams(self, data):
 70 |         self.hparams.update(data.hparams)
 71 | 
 72 |     @property
 73 |     def example_input_array(self):
 74 |         common = torch.randn(16, self.in_channels, self.seq_len)
 75 | 
 76 |         return common
 77 | 
 78 |     def configure_optimizers(self):
 79 |         param_groups = [
 80 |             {"params": self.encoder.parameters()},
 81 |             {"params": self.regressor.parameters()},
 82 |         ]
 83 |         if self.optim_type == "adam":
 84 |             return torch.optim.Adam(param_groups, lr=self.lr)
 85 |         else:
 86 |             return torch.optim.SGD(
 87 |                 param_groups, lr=self.lr, momentum=0.9, weight_decay=0.01
 88 |             )
 89 | 
 90 |     def forward(self, inputs):
 91 |         latent_code = self.encoder(inputs)
 92 |         prediction = self.regressor(latent_code)
 93 | 
 94 |         return prediction
 95 | 
 96 |     def training_step(self, batch, batch_idx):
 97 |         source, source_labels = batch
 98 |         predictions = self(source)
 99 |         loss = self.criterion_regression(predictions, source_labels)
100 | 
101 |         self.log("train/regression_loss", loss)
102 | 
103 |         return loss
104 | 
105 |     def on_validation_epoch_start(self):
106 |         self._reset_all_metrics()
107 | 
108 |     def on_test_epoch_start(self):
109 |         self._reset_all_metrics()
110 | 
111 |     def _reset_all_metrics(self):
112 |         self.embedding_metric.reset()
113 |         for metric in self.regression_metrics.values():
114 |             metric.reset()
115 | 
116 |     def validation_step(self, batch, batch_idx):
117 |         self._evaluate(batch, metric_id=1)
118 | 
119 |     def test_step(self, batch, batch_idx, dataloader_idx):
120 |         self._evaluate(batch, metric_id=dataloader_idx + 1)
121 | 
122 |     def _evaluate(self, batch, metric_id):
123 |         features, labels = batch
124 |         predictions = self(features)
125 |         self.regression_metrics[metric_id].update(predictions, labels)
126 |         if self.record_embeddings:
127 |             latent_code = self.encoder(features)
128 |             self.embedding_metric.update(latent_code, torch.zeros_like(labels), labels)
129 | 
130 |     def validation_epoch_end(self, outputs):
131 |         self.log("val/regression_loss", self.regression_metrics[1].compute())
132 |         if self.record_embeddings:
133 |             fig_class, fig_rul = self.embedding_metric.compute()
134 |             self.logger.log_figure("val/embeddings_class", fig_class, self.global_step)
135 |             self.logger.log_figure("val/embeddings_rul", fig_rul, self.global_step)
136 |             self.embedding_metric.reset()
137 | 
138 |     def test_epoch_end(self, outputs):
139 |         for fd, metric in self.regression_metrics.items():
140 |             self.log(f"test/regression_loss_fd{fd}", metric.compute())
141 | 


--------------------------------------------------------------------------------
/src/lightning/loggers.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | import tempfile
  3 | from typing import Optional
  4 | 
  5 | import pytorch_lightning.loggers as loggers
  6 | from pytorch_lightning.callbacks import ModelCheckpoint
  7 | 
  8 | ExperimentNaming = {1: "one", 2: "two", 3: "three", 4: "four"}
  9 | 
 10 | 
 11 | def baseline_experiment_name(source):
 12 |     assert source in ExperimentNaming, f"Unknown FD number {source}."
 13 |     return f"cmapss_{ExperimentNaming[source]}_baseline"
 14 | 
 15 | 
 16 | def transfer_experiment_name(source, target):
 17 |     assert source in ExperimentNaming, f"Unknown source FD number {source}."
 18 |     assert target in ExperimentNaming, f"Unknown target FD number {target}."
 19 |     return f"{ExperimentNaming[source]}2{ExperimentNaming[target]}"
 20 | 
 21 | 
 22 | def pretraining_experiment_name(source, target):
 23 |     if target is None:
 24 |         return pretraining_baseline_experiment_name(source)
 25 |     else:
 26 |         return pretraining_adaption_experiment_name(source, target)
 27 | 
 28 | 
 29 | def pretraining_baseline_experiment_name(dataset):
 30 |     assert dataset in ExperimentNaming, f"Unknown dataset FD number {dataset}."
 31 | 
 32 |     return f"pretraining_{ExperimentNaming[dataset]}"
 33 | 
 34 | 
 35 | def pretraining_adaption_experiment_name(source, target):
 36 |     assert source in ExperimentNaming, f"Unknown source FD number {source}."
 37 |     assert target in ExperimentNaming, f"Unknown target FD number {target}."
 38 | 
 39 |     return f"pretraining_{ExperimentNaming[source]}&{ExperimentNaming[target]}"
 40 | 
 41 | 
 42 | def transfer_hyperopt_name(source, target, percent_broken):
 43 |     assert source in ExperimentNaming, f"Unknown source FD number {source}."
 44 |     assert target in ExperimentNaming, f"Unknown target FD number {target}."
 45 |     assert 0 <= percent_broken <= 1, f"Invalid percent broken {percent_broken}."
 46 | 
 47 |     return f"hyperopt_{ExperimentNaming[source]}&{ExperimentNaming[target]}@{percent_broken:.2f}"
 48 | 
 49 | 
 50 | def pretraining_hyperopt_name(source, target, percent_broken):
 51 |     assert source in ExperimentNaming, f"Unknown source FD number {source}."
 52 |     assert target in ExperimentNaming, f"Unknown target FD number {target}."
 53 |     assert 0 <= percent_broken <= 1, f"Invalid percent broken {percent_broken}."
 54 | 
 55 |     return f"pre_hyperopt_{ExperimentNaming[source]}&{ExperimentNaming[target]}@{percent_broken:.2f}"
 56 | 
 57 | 
 58 | def semi_supervised_hyperopt_name(source, percent_broken):
 59 |     assert source in ExperimentNaming, f"Unknown source FD number {source}."
 60 |     assert 0 <= percent_broken <= 1, f"Invalid percent broken {percent_broken}."
 61 | 
 62 |     return f"pre_hyperopt_{ExperimentNaming[source]}@{percent_broken:.2f}"
 63 | 
 64 | 
 65 | class MLTBLogger(loggers.LoggerCollection):
 66 |     """Combined MlFlow and Tensorboard logger that saves models as MlFlow artifacts."""
 67 | 
 68 |     def __init__(self, log_dir, experiment_name, tag=None, tensorboard_struct=None):
 69 |         """
 70 |         This logger combines a MlFlow and Tensorboard logger.
 71 |         It creates a directory (mlruns/tensorboard) for each logger in log_dir.
 72 | 
 73 |         If a tensorboard_struct dict is provided, it is used to create additional
 74 |         sub-directories for tensorboard to get a better overview over the runs.
 75 | 
 76 |         The difference to a simple LoggerCollection is that the save dir points
 77 |         to the artifact path of the MlFlow run. This way the model is logged as
 78 |         a MlFlow artifact.
 79 | 
 80 |         :param log_dir: directory to put the mlruns and tensorboard directories
 81 |         :param experiment_name: name for the experiment
 82 |         :param tensorboard_struct: dictionary containing information to refine the tensorboard directory structure
 83 |         """
 84 |         tensorboard_path = os.path.join(log_dir, "tensorboard", experiment_name)
 85 |         sub_dirs = self._dirs_from_dict(tensorboard_struct)
 86 |         self._tf_logger = loggers.TensorBoardLogger(tensorboard_path, name=sub_dirs)
 87 | 
 88 |         mlflow_path = "file:" + os.path.normpath(os.path.join(log_dir, "mlruns"))
 89 |         self._mlflow_logger = loggers.MLFlowLogger(
 90 |             experiment_name, tracking_uri=mlflow_path, tags=self._build_tags(tag)
 91 |         )
 92 | 
 93 |         super().__init__([self._tf_logger, self._mlflow_logger])
 94 | 
 95 |     @staticmethod
 96 |     def _dirs_from_dict(tensorboard_struct):
 97 |         if tensorboard_struct is not None:
 98 |             dirs = []
 99 |             for key, value in tensorboard_struct.items():
100 |                 if isinstance(value, float) and value >= 0.1:
101 |                     dirs.append(f"{value:.1f}{key}")
102 |                 elif isinstance(value, str):
103 |                     dirs.append(f"{key}-{value}")
104 |                 else:
105 |                     dirs.append(f"{value}{key}")
106 |             dirs = os.path.join(*dirs)
107 |         else:
108 |             dirs = ""
109 | 
110 |         return dirs
111 | 
112 |     def _build_tags(self, tag):
113 |         if tag is not None:
114 |             tags = {"version": tag}
115 |         else:
116 |             tags = None
117 | 
118 |         return tags
119 | 
120 |     @property
121 |     def name(self) -> str:
122 |         return self._mlflow_logger.name
123 | 
124 |     @property
125 |     def version(self) -> str:
126 |         return os.path.join(self._mlflow_logger.version, "artifacts")
127 | 
128 |     @property
129 |     def save_dir(self):
130 |         return self._mlflow_logger.save_dir
131 | 
132 |     @property
133 |     def checkpoint_path(self):
134 |         return os.path.join(self.save_dir, self.name, self.version, "checkpoints")
135 | 
136 |     @property
137 |     def tf_experiment(self):
138 |         return self._tf_logger.experiment
139 | 
140 |     @property
141 |     def mlflow_experiment(self):
142 |         return self._mlflow_logger.experiment
143 | 
144 |     @property
145 |     def run_id(self):
146 |         return self._mlflow_logger.run_id
147 | 
148 |     def log_figure(self, tag, figure, step):
149 |         tag_tail, tag_head = os.path.split(tag)
150 |         _, tmp_file = tempfile.mkstemp(".pdf", f"{tag_head}_{step:05}_")
151 |         figure.savefig(tmp_file)
152 |         self.mlflow_experiment.log_artifact(
153 |             self._mlflow_logger.run_id, tmp_file, tag_tail
154 |         )
155 |         self.tf_experiment.add_figure(tag, figure, step)
156 | 
157 | 
158 | class MinEpochModelCheckpoint(ModelCheckpoint):
159 |     """Checkpoints models only after training for a minimum of epochs."""
160 | 
161 |     def __init__(
162 |         self,
163 |         filepath: Optional[str] = None,
164 |         monitor: Optional[str] = None,
165 |         verbose: bool = False,
166 |         save_last: Optional[bool] = None,
167 |         save_top_k: Optional[int] = None,
168 |         save_weights_only: bool = False,
169 |         mode: str = "auto",
170 |         period: int = 1,
171 |         prefix: str = "",
172 |         min_epochs_before_saving: int = 0,
173 |     ):
174 |         super().__init__(
175 |             filepath,
176 |             monitor,
177 |             verbose,
178 |             save_last,
179 |             save_top_k,
180 |             save_weights_only,
181 |             mode,
182 |             period,
183 |             prefix,
184 |         )
185 | 
186 |         self.min_epochs_before_saving = min_epochs_before_saving
187 | 
188 |     def save_checkpoint(self, trainer, pl_module):
189 |         if trainer.current_epoch > self.min_epochs_before_saving:
190 |             super().save_checkpoint(trainer, pl_module)
191 | 


--------------------------------------------------------------------------------
/src/lightning/metrics.py:
--------------------------------------------------------------------------------
  1 | import matplotlib.colors as mplcolors
  2 | import matplotlib.pyplot as plt
  3 | import pytorch_lightning as pl
  4 | import torch
  5 | import umap
  6 | from torch import nn as nn
  7 | 
  8 | 
  9 | class EmbeddingViz(pl.metrics.Metric):
 10 |     def __init__(self, num_elements, embedding_size, combined=True):
 11 |         super().__init__()
 12 | 
 13 |         self.num_elements = num_elements
 14 |         self.embedding_size = embedding_size
 15 |         self.combined = combined
 16 | 
 17 |         self.add_state(
 18 |             "embeddings",
 19 |             default=torch.zeros(num_elements, embedding_size),
 20 |             dist_reduce_fx=None,
 21 |         )
 22 |         self.add_state("labels", default=torch.zeros(num_elements), dist_reduce_fx=None)
 23 |         self.add_state("ruls", default=torch.zeros(num_elements), dist_reduce_fx=None)
 24 |         self.add_state("sample_counter", default=torch.tensor(0), dist_reduce_fx=None)
 25 | 
 26 |         self.class_cm = plt.get_cmap("tab10")
 27 |         self.rul_cm = plt.get_cmap("viridis")
 28 | 
 29 |     def update(self, embeddings, labels, ruls):
 30 |         """
 31 |         Add embedding, labels and RUL to the metric.
 32 | 
 33 |         :param embeddings: hidden layer embedding of the data points
 34 |         :param labels: domain labels with 0 being source and 1 being target
 35 |         :param ruls: Remaining Useful Lifetime values of the data points
 36 |         """
 37 |         start = self.sample_counter
 38 |         end = self.sample_counter + embeddings.shape[0]
 39 |         self.sample_counter = end
 40 | 
 41 |         self.embeddings[start:end] = embeddings
 42 |         self.labels[start:end] = labels
 43 |         self.ruls[start:end] = ruls
 44 | 
 45 |     def compute(self):
 46 |         """Compute UMAP and plot points to 2d scatter plot."""
 47 |         logged_embeddings = self.embeddings[: self.sample_counter].detach().cpu().numpy()
 48 |         logged_labels = self.labels[: self.sample_counter].detach().cpu().int()
 49 |         logged_ruls = self.ruls[: self.sample_counter].detach().cpu()
 50 |         viz_embeddings = umap.UMAP(random_state=42).fit_transform(logged_embeddings)
 51 | 
 52 |         if self.combined:
 53 |             fig, (ax_class, ax_rul) = plt.subplots(
 54 |                 1, 2, sharex="all", sharey="all", figsize=(20, 10), dpi=300
 55 |             )
 56 |         else:
 57 |             fig_class = plt.figure(figsize=(10, 10), dpi=300)
 58 |             ax_class = fig_class.gca()
 59 |             fig_rul = plt.figure(figsize=(10, 10), dpi=300)
 60 |             ax_rul = fig_rul.gca()
 61 |             fig = [fig_class, fig_rul]
 62 | 
 63 |         self._scatter_class(ax_class, viz_embeddings, logged_labels)
 64 |         self._scatter_rul(ax_rul, viz_embeddings, logged_ruls)
 65 | 
 66 |         return fig
 67 | 
 68 |     def _scatter_class(self, ax, embeddings, labels):
 69 |         class_colors = [self.class_cm.colors[c] for c in labels]
 70 |         ax.scatter(
 71 |             embeddings[:, 0],
 72 |             embeddings[:, 1],
 73 |             c=class_colors,
 74 |             alpha=0.4,
 75 |             edgecolors="none",
 76 |             s=[4],
 77 |             rasterized=True,
 78 |         )
 79 |         ax.legend(["Source", "Target"], loc="lower left")
 80 | 
 81 |     def _scatter_rul(self, ax, embeddings, ruls):
 82 |         ax.scatter(
 83 |             embeddings[:, 0],
 84 |             embeddings[:, 1],
 85 |             c=ruls,
 86 |             alpha=0.4,
 87 |             edgecolors="none",
 88 |             s=[4],
 89 |             rasterized=True,
 90 |         )
 91 |         color_bar = plt.cm.ScalarMappable(
 92 |             mplcolors.Normalize(ruls.min(), ruls.max()), self.rul_cm
 93 |         )
 94 |         color_bar.set_array(ruls)
 95 |         plt.colorbar(color_bar)
 96 | 
 97 | 
 98 | class RMSELoss(pl.metrics.Metric):
 99 |     def __init__(self, num_elements: int = 1000):
100 |         super().__init__()
101 | 
102 |         self.mse = nn.MSELoss()
103 | 
104 |         self.add_state("losses", default=torch.zeros(num_elements), dist_reduce_fx=None)
105 |         self.add_state("sizes", default=torch.zeros(num_elements), dist_reduce_fx=None)
106 |         self.add_state("sample_counter", default=torch.tensor(0), dist_reduce_fx=None)
107 | 
108 |     def update(self, inputs: torch.Tensor, targets: torch.Tensor):
109 |         summed_square = nn.functional.mse_loss(inputs, targets, reduction="sum")
110 |         batch_size = inputs.shape[0]
111 | 
112 |         self.losses[self.sample_counter] = summed_square
113 |         self.sizes[self.sample_counter] = batch_size
114 |         self.sample_counter += 1
115 | 
116 |     def compute(self) -> torch.Tensor:
117 |         if self.sample_counter == 0:
118 |             raise RuntimeError("RMSE metric was not used. Computation impossible.")
119 |         summed_squares = self.losses[: self.sample_counter]
120 |         batch_sizes = self.sizes[: self.sample_counter]
121 |         rmse = torch.sqrt(summed_squares.sum() / batch_sizes.sum())
122 | 
123 |         return rmse
124 | 
125 |     def forward(self, inputs, targets):
126 |         return torch.sqrt(self.mse(inputs, targets))
127 | 
128 | 
129 | class SimpleMetric(pl.metrics.Metric):
130 |     def __init__(self, num_elements: int = 1000, reduction="mean"):
131 |         super().__init__()
132 | 
133 |         if reduction not in ["sum", "mean"]:
134 |             raise ValueError(f"Unsupported reduction {reduction}")
135 |         self.reduction = reduction
136 | 
137 |         self.add_state("losses", default=torch.zeros(num_elements), dist_reduce_fx=None)
138 |         self.add_state("sizes", default=torch.zeros(num_elements), dist_reduce_fx=None)
139 |         self.add_state("sample_counter", default=torch.tensor(0), dist_reduce_fx=None)
140 | 
141 |     def update(self, loss: torch.Tensor, batch_size: int):
142 |         self.losses[self.sample_counter] = loss
143 |         self.sizes[self.sample_counter] = batch_size
144 |         self.sample_counter += 1
145 | 
146 |     def compute(self) -> torch.Tensor:
147 |         if self.sample_counter == 0:
148 |             raise RuntimeError("RMSE metric was not used. Computation impossible.")
149 |         if self.reduction == "mean":
150 |             loss = self._weighted_mean()
151 |         else:
152 |             loss = self._sum()
153 | 
154 |         return loss
155 | 
156 |     def _weighted_mean(self):
157 |         weights = self.sizes[: self.sample_counter]
158 |         weights = weights / weights.sum()
159 |         loss = self.losses[: self.sample_counter]
160 |         loss = torch.sum(loss * weights)
161 | 
162 |         return loss
163 | 
164 |     def _sum(self):
165 |         loss = self.losses[: self.sample_counter]
166 |         loss = torch.sum(loss)
167 | 
168 |         return loss
169 | 
170 | 
171 | class RULScore:
172 |     def __init__(self, pos_factor=10, neg_factor=-13):
173 |         self.pos_factor = pos_factor
174 |         self.neg_factor = neg_factor
175 | 
176 |     def __call__(self, inputs, targets):
177 |         dist = inputs - targets
178 |         for i, d in enumerate(dist):
179 |             dist[i] = (d / self.neg_factor) if d < 0 else (d / self.pos_factor)
180 |         dist = torch.exp(dist) - 1
181 |         score = dist.sum()
182 | 
183 |         return score
184 | 


--------------------------------------------------------------------------------
/src/lightning/mixins.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | 
 3 | 
 4 | class DataHparamsMixin:
 5 |     def add_data_hparams(self, data):
 6 |         self.hparams.update(data.hparams)
 7 | 
 8 | 
 9 | class LoadEncoderMixin:
10 |     def load_encoder(self, checkpoint_path, load_disc=False):
11 |         checkpoint = torch.load(checkpoint_path, map_location="cpu")
12 | 
13 |         encoder_state = self._extract_state_dict(checkpoint, model="encoder")
14 |         self.encoder.load_state_dict(encoder_state)
15 |         if load_disc:
16 |             disc_state = self._extract_state_dict(checkpoint, model="domain_disc")
17 |             self.domain_disc.load_state_dict(disc_state)
18 | 
19 |         self.hparams["pretrained_checkpoint"] = checkpoint_path
20 |         self.encoder.norm_outputs = True
21 | 
22 |     def _extract_state_dict(self, checkpoint, model):
23 |         encoder_state = {
24 |             n.replace(model + ".", ""): weight
25 |             for n, weight in checkpoint["state_dict"].items()
26 |             if n.startswith(model)
27 |         }
28 | 
29 |         return encoder_state
30 | 
31 |     def load_from_rbm(self, rbm):
32 |         first_conv = self.encoder.layers[0]
33 |         first_bn = self.encoder.layers[1]
34 |         first_conv.weight.data = rbm.interaction.module.weight.data
35 |         first_bn.bias.data = rbm.hidden.bias.data.squeeze()
36 | 


--------------------------------------------------------------------------------
/src/lightning/pretraining.py:
--------------------------------------------------------------------------------
  1 | import pytorch_lightning as pl
  2 | import torch
  3 | import torch.nn as nn
  4 | 
  5 | from lightning import metrics
  6 | from lightning.mixins import DataHparamsMixin
  7 | from models import networks
  8 | 
  9 | 
 10 | class UnsupervisedPretraining(pl.LightningModule, DataHparamsMixin):
 11 |     def __init__(
 12 |         self,
 13 |         in_channels,
 14 |         seq_len,
 15 |         num_layers,
 16 |         kernel_size,
 17 |         base_filters,
 18 |         latent_dim,
 19 |         dropout,
 20 |         domain_tradeoff,
 21 |         domain_disc_dim,
 22 |         num_disc_layers,
 23 |         lr,
 24 |         weight_decay,
 25 |         record_embeddings=False,
 26 |         encoder="cnn",
 27 |     ):
 28 |         super().__init__()
 29 | 
 30 |         self.in_channels = in_channels
 31 |         self.seq_len = seq_len
 32 |         self.num_layers = num_layers
 33 |         self.kernel_size = kernel_size
 34 |         self.base_filters = base_filters
 35 |         self.latent_dim = latent_dim
 36 |         self.dropout = dropout
 37 |         self.domain_tradeoff = domain_tradeoff
 38 |         self.domain_disc_dim = domain_disc_dim
 39 |         self.num_disc_layers = num_disc_layers
 40 |         self.lr = lr
 41 |         self.weight_decay = weight_decay
 42 |         self.record_embeddings = record_embeddings
 43 |         self.encoder_type = encoder
 44 | 
 45 |         self.encoder = self._get_encoder()
 46 |         self.domain_disc = None
 47 | 
 48 |         self.criterion_regression = nn.MSELoss()
 49 |         self.criterion_domain = nn.BCEWithLogitsLoss()
 50 | 
 51 |         self.embedding_metric = metrics.EmbeddingViz(40000, self.latent_dim)
 52 |         self.regression_metric = metrics.SimpleMetric(40000)
 53 |         self.domain_metric = metrics.SimpleMetric(40000)
 54 | 
 55 |         self.save_hyperparameters()
 56 |         self.hparams["mode"] = "metric"
 57 | 
 58 |     def _get_encoder(self):
 59 |         if self.encoder_type == "cnn":
 60 |             encoder = networks.Encoder(
 61 |                 self.in_channels,
 62 |                 self.base_filters,
 63 |                 self.kernel_size,
 64 |                 self.num_layers,
 65 |                 self.latent_dim,
 66 |                 self.seq_len,
 67 |                 self.dropout,
 68 |                 norm_outputs=True,
 69 |             )
 70 |         else:
 71 |             raise ValueError(f"Unknown encoder type {self.encoder_type}. Use 'cnn'.")
 72 |         return encoder
 73 | 
 74 |     def configure_optimizers(self):
 75 |         return torch.optim.Adam(
 76 |             self.parameters(), lr=self.lr, weight_decay=self.weight_decay
 77 |         )
 78 | 
 79 |     def forward(self, anchors, queries):
 80 |         anchor_embeddings, query_embeddings = self._get_anchor_query_embeddings(
 81 |             anchors, queries
 82 |         )
 83 |         pred_distances = self._pairwise_distance(anchor_embeddings, query_embeddings)
 84 | 
 85 |         return pred_distances
 86 | 
 87 |     def training_step(self, batch, batch_idx):
 88 |         regression_loss, domain_loss = self._get_losses(batch)
 89 |         loss = regression_loss + self.domain_tradeoff * domain_loss
 90 |         self.log("train/loss", loss)
 91 |         self.log("train/regression_loss", regression_loss)
 92 |         self.log("train/domain_loss", domain_loss)
 93 | 
 94 |         return loss
 95 | 
 96 |     def on_validation_epoch_start(self):
 97 |         self._reset_all_metrics()
 98 | 
 99 |     def _reset_all_metrics(self):
100 |         self.embedding_metric.reset()
101 |         self.regression_metric.reset()
102 |         self.domain_metric.reset()
103 | 
104 |     def validation_step(self, batch, batch_idx, dataloader_idx):
105 |         if dataloader_idx == 0:
106 |             regression_loss, domain_loss = self._get_losses(batch)
107 |             batch_size = batch[0].shape[0]
108 |             self.regression_metric.update(regression_loss, batch_size)
109 |             self.domain_metric.update(domain_loss, batch_size)
110 |         elif self.record_embeddings:
111 |             self._record_embeddings(batch, dataloader_idx)
112 | 
113 |     def _record_embeddings(self, batch, dataloader_idx):
114 |         features, labels = batch
115 |         embedding = self.encoder(features)
116 |         domain_labels = torch.full_like(
117 |             labels, fill_value=(2 - dataloader_idx), dtype=torch.int
118 |         )
119 |         self.embedding_metric.update(embedding, domain_labels, labels)
120 | 
121 |     def validation_epoch_end(self, validation_step_outputs):
122 |         if self.record_embeddings:
123 |             embedding_fig = self.embedding_metric.compute()
124 |             self.logger.log_figure("val/embeddings", embedding_fig, self.global_step)
125 | 
126 |         regression_loss = self.regression_metric.compute()
127 |         domain_loss = self.domain_metric.compute()
128 | 
129 |         self.log("val/regression_loss", regression_loss)
130 |         self.log("val/domain_loss", domain_loss)
131 |         self.log("val/checkpoint_score", regression_loss - 0.1 * domain_loss)
132 | 
133 |     def _get_losses(self, batch):
134 |         anchors, queries, true_distances, domain_labels = batch
135 |         anchor_embeddings, query_embeddings = self._get_anchor_query_embeddings(
136 |             anchors, queries
137 |         )
138 |         pred_distances = self._pairwise_distance(anchor_embeddings, query_embeddings)
139 |         regression_loss = self.criterion_regression(pred_distances, true_distances)
140 | 
141 |         if self.domain_tradeoff > 0:
142 |             domain_pred = self.domain_disc(anchor_embeddings)
143 |             domain_loss = self.criterion_domain(domain_pred, domain_labels)
144 |         else:
145 |             domain_loss = 0
146 | 
147 |         return regression_loss, domain_loss
148 | 
149 |     def _pairwise_distance(self, anchor_embeddings, query_embeddings):
150 |         return torch.pairwise_distance(anchor_embeddings, query_embeddings, eps=1e-8)
151 | 
152 |     def _get_anchor_query_embeddings(self, anchors, queries):
153 |         batch_size = anchors.shape[0]
154 |         combined = torch.cat([anchors, queries])
155 |         embeddings = self.encoder(combined)
156 |         anchor_embeddings, query_embeddings = torch.split(embeddings, batch_size)
157 | 
158 |         return anchor_embeddings, query_embeddings
159 | 


--------------------------------------------------------------------------------
/src/models/layers.py:
--------------------------------------------------------------------------------
  1 | from typing import Union
  2 | 
  3 | import torch
  4 | from pytorch_probgraph import GaussianLayer, UnitLayer
  5 | from torch import nn as nn
  6 | 
  7 | 
  8 | class DeFlatten(nn.Module):
  9 |     def __init__(self, seq_len, num_channels):
 10 |         super().__init__()
 11 | 
 12 |         self.seq_len = seq_len
 13 |         self.num_channels = num_channels
 14 | 
 15 |     def forward(self, inputs):
 16 |         return inputs.view(-1, self.num_channels, self.seq_len)
 17 | 
 18 | 
 19 | class _GradientReverse(torch.autograd.Function):
 20 |     """Gradient reversal forward and backward definitions."""
 21 | 
 22 |     @staticmethod
 23 |     def forward(ctx, inputs, **kwargs):
 24 |         """Forward pass as identity mapping."""
 25 |         return inputs
 26 | 
 27 |     @staticmethod
 28 |     def backward(ctx, grad):
 29 |         """Backward pass as negative of gradient."""
 30 |         return -grad
 31 | 
 32 | 
 33 | def gradient_reversal(x):
 34 |     """Perform gradient reversal on input."""
 35 |     return _GradientReverse.apply(x)
 36 | 
 37 | 
 38 | class GradientReversalLayer(nn.Module):
 39 |     """Module for gradient reversal."""
 40 | 
 41 |     def forward(self, inputs):
 42 |         """Perform forward pass of gradient reversal."""
 43 |         return gradient_reversal(inputs)
 44 | 
 45 | 
 46 | class RectifiedLinearLayer(UnitLayer):
 47 |     """
 48 |     A UnitLayer of rectified linear units.
 49 |     """
 50 | 
 51 |     def __init__(self, bias: torch.Tensor):
 52 |         """
 53 |         :param bias: Bias for the ReLU unit.
 54 |         """
 55 |         super().__init__()
 56 |         self.register_parameter("bias", torch.nn.Parameter(bias))
 57 | 
 58 |     def mean_cond(
 59 |         self, interaction: Union[torch.Tensor, None] = None, N: int = 1
 60 |     ) -> torch.Tensor:
 61 |         return nn.functional.relu(interaction + self.bias)
 62 | 
 63 |     def sample_cond(
 64 |         self, interaction: Union[torch.Tensor, None] = None, N: int = 1
 65 |     ) -> torch.Tensor:
 66 |         interaction = interaction + self.bias
 67 |         return nn.functional.relu(interaction + torch.sigmoid(interaction))
 68 | 
 69 |     def transform(self, input: torch.Tensor) -> torch.Tensor:
 70 |         return input
 71 | 
 72 |     def transform_invert(self, transformed_input: torch.Tensor) -> torch.Tensor:
 73 |         return transformed_input
 74 | 
 75 |     def logprob_cond(
 76 |         self, input: torch.Tensor, interaction: Union[torch.Tensor, float] = 0.0
 77 |     ) -> torch.Tensor:
 78 |         raise NotImplementedError
 79 | 
 80 |     def logprob_joint(self, input: torch.Tensor) -> torch.Tensor:
 81 |         return torch.sum(input * self.bias, dim=list(range(1, len(input.shape))))
 82 | 
 83 |     def free_energy(self, interaction: torch.Tensor) -> torch.Tensor:
 84 |         raise NotImplementedError
 85 | 
 86 |     def backward(
 87 |         self, input: torch.Tensor, factor: Union[torch.Tensor, float] = 1.0
 88 |     ) -> None:
 89 |         if self.bias.requires_grad:
 90 |             grad_bias = (factor * input).mean(dim=0, keepdim=True)
 91 |             grad_bias = grad_bias.mean(dim=-1, keepdim=True)
 92 |             self.bias.backward(grad_bias.detach())
 93 | 
 94 | 
 95 | class GaussianSequenceLayer(GaussianLayer):
 96 |     def backward(
 97 |         self, input: torch.Tensor, factor: Union[torch.Tensor, float] = 1.0
 98 |     ) -> None:
 99 |         if self.bias.requires_grad:
100 |             grad_bias = (factor * input).sum(dim=0, keepdim=True) / input.shape[0]
101 |             grad_bias = grad_bias.sum(dim=-1, keepdim=True) / input.shape[-1]
102 |             self.bias.backward(grad_bias.detach())
103 |         if self.logsigma.requires_grad:
104 |             var = (input ** 2).sum(dim=0, keepdim=True) / input.shape[0]
105 |             grad_logsigma = factor * var * torch.exp(-2 * self.logsigma) - 1.0
106 |             self.logsigma.backward(grad_logsigma.detach())
107 | 


--------------------------------------------------------------------------------
/src/models/networks.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | import torch.nn as nn
  3 | 
  4 | from models import layers
  5 | 
  6 | 
  7 | class Encoder(nn.Module):
  8 |     def __init__(
  9 |         self,
 10 |         in_channels,
 11 |         base_filters,
 12 |         kernel_size,
 13 |         num_layers,
 14 |         latent_dim,
 15 |         seq_len,
 16 |         dropout,
 17 |         norm_outputs,
 18 |     ):
 19 |         super().__init__()
 20 | 
 21 |         self.in_channels = in_channels
 22 |         self.base_filters = base_filters
 23 |         self.kernel_size = kernel_size
 24 |         self.num_layers = num_layers
 25 |         self.latent_dim = latent_dim
 26 |         self.seq_len = seq_len
 27 |         self.dropout = dropout
 28 |         self.norm_outputs = norm_outputs
 29 | 
 30 |         self.layers = self._build_encoder()
 31 | 
 32 |     @property
 33 |     def padding(self):
 34 |         left = self.kernel_size - 1
 35 |         right = 0
 36 | 
 37 |         return left, right
 38 | 
 39 |     def _build_encoder(self):
 40 |         sequence = [
 41 |             nn.Conv1d(self.in_channels, self.base_filters, self.kernel_size, bias=False),
 42 |             nn.BatchNorm1d(self.base_filters),
 43 |             nn.ReLU(True),
 44 |         ]
 45 |         for i in range(1, self.num_layers):
 46 |             in_filters = min(i * self.base_filters, 64)
 47 |             out_filters = min((i + 1) * self.base_filters, 64)
 48 |             use_padding = i % 2 == 1
 49 |             sequence.extend(self._build_conv_layer(in_filters, out_filters, use_padding))
 50 | 
 51 |         cut_off = self.num_layers // 2 * (self.kernel_size - (self.kernel_size % 2))
 52 |         flat_dim = (self.seq_len - cut_off) * min(self.num_layers * self.base_filters, 64)
 53 |         sequence.extend([nn.Flatten(), nn.Linear(flat_dim, self.latent_dim)])
 54 | 
 55 |         return nn.Sequential(*sequence)
 56 | 
 57 |     def _build_conv_layer(self, in_filters, out_filters, use_padding):
 58 |         layer = []
 59 |         if use_padding:
 60 |             layer.append(nn.ConstantPad1d(self.padding, 0.0))
 61 |         layer.extend(
 62 |             [
 63 |                 nn.Conv1d(in_filters, out_filters, self.kernel_size, bias=False),
 64 |                 nn.BatchNorm1d(out_filters),
 65 |                 nn.ReLU(True),
 66 |                 nn.Dropout2d(p=self.dropout),
 67 |             ]
 68 |         )
 69 | 
 70 |         return layer
 71 | 
 72 |     def forward(self, inputs):
 73 |         outputs = self.layers(inputs)
 74 |         if self.norm_outputs:
 75 |             outputs = outputs / torch.norm(outputs, dim=1, keepdim=True)
 76 | 
 77 |         return outputs
 78 | 
 79 | 
 80 | class Decoder(nn.Module):
 81 |     def __init__(
 82 |         self,
 83 |         in_channels,
 84 |         base_filters,
 85 |         kernel_size,
 86 |         num_layers,
 87 |         latent_dim,
 88 |         seq_len,
 89 |         dropout,
 90 |     ):
 91 |         super().__init__()
 92 | 
 93 |         self.in_channels = in_channels
 94 |         self.base_filters = base_filters
 95 |         self.kernel_size = kernel_size
 96 |         self.num_layers = num_layers
 97 |         self.latent_dim = latent_dim
 98 |         self.seq_len = seq_len
 99 |         self.dropout = dropout
100 | 
101 |         self.layers = self._build_decoder()
102 | 
103 |     @property
104 |     def padding(self):
105 |         left = self.kernel_size - 1
106 |         right = 0
107 | 
108 |         return left, right
109 | 
110 |     def _build_decoder(self):
111 |         cut_off = self.num_layers // 2 * (self.kernel_size - (self.kernel_size % 2))
112 |         max_filters = min(self.num_layers * self.base_filters, 64)
113 |         reduced_seq_len = self.seq_len - cut_off
114 |         flat_dim = reduced_seq_len * max_filters
115 | 
116 |         sequence = [
117 |             nn.Linear(self.latent_dim, flat_dim),
118 |             nn.BatchNorm1d(flat_dim),
119 |             nn.ReLU(True),
120 |             layers.DeFlatten(reduced_seq_len, max_filters),
121 |         ]
122 |         for i in range(self.num_layers - 1, 0, -1):
123 |             in_filters = min((i + 1) * self.base_filters, 64)
124 |             out_filters = min(i * self.base_filters, 64)
125 |             use_padding = i % 2 == 1
126 |             sequence.extend(self._build_conv_layer(in_filters, out_filters, use_padding))
127 | 
128 |         sequence.extend(
129 |             [
130 |                 nn.ConvTranspose1d(self.base_filters, self.in_channels, self.kernel_size),
131 |                 nn.Tanh(),
132 |             ]
133 |         )
134 | 
135 |         return nn.Sequential(*sequence)
136 | 
137 |     def _build_conv_layer(self, in_filters, out_filters, use_padding):
138 |         layer = []
139 |         if use_padding:
140 |             layer.extend(
141 |                 [
142 |                     nn.ConstantPad1d(self.padding, 0.0),
143 |                     nn.Conv1d(in_filters, out_filters, self.kernel_size, bias=False),
144 |                 ]
145 |             )
146 |         else:
147 |             layer.append(
148 |                 nn.ConvTranspose1d(in_filters, out_filters, self.kernel_size, bias=False)
149 |             )
150 |         layer.extend(
151 |             [
152 |                 nn.BatchNorm1d(out_filters),
153 |                 nn.ReLU(True),
154 |                 nn.Dropout2d(p=self.dropout),
155 |             ]
156 |         )
157 | 
158 |         return layer
159 | 
160 |     def forward(self, inputs):
161 |         outputs = inputs
162 |         for m in self.layers:
163 |             outputs = m(outputs)
164 | 
165 |         return outputs
166 | 
167 | 
168 | class Regressor(nn.Module):
169 |     def __init__(self, latent_dim):
170 |         super().__init__()
171 | 
172 |         self.latent_dim = latent_dim
173 | 
174 |         self.layers = self._build_regressor()
175 | 
176 |     def _build_regressor(self):
177 |         classifier = nn.Sequential(
178 |             nn.BatchNorm1d(self.latent_dim),
179 |             nn.ReLU(True),
180 |             nn.Linear(self.latent_dim, 1),
181 |         )
182 | 
183 |         return classifier
184 | 
185 |     def forward(self, inputs):
186 |         return self.layers(inputs).squeeze(1)
187 | 


--------------------------------------------------------------------------------
/src/reproduce_baseline.py:
--------------------------------------------------------------------------------
 1 | import os
 2 | from datetime import datetime
 3 | import random
 4 | 
 5 | import ray
 6 | 
 7 | import building
 8 | from run_semi_supervised import run
 9 | 
10 | 
11 | script_path = os.path.dirname(__file__)
12 | config_root = os.path.join(script_path, "..", "configs")
13 | 
14 | 
15 | def reproduce(base_version, percent_fails, encoder, master_seed):
16 |     ray.init()
17 | 
18 |     if base_version is None:
19 |         base_version = datetime.now().timestamp()
20 |     error_log = []
21 |     fds = [1, 2, 3, 4]
22 |     random.seed(master_seed)
23 |     seeds = [
24 |         [random.randint(0, 99999999) for _ in range(len(percent_fails))]
25 |         for _ in range(len(fds))
26 |     ]
27 | 
28 |     for num_fd, fd in enumerate(fds):
29 |         arch_config_path = os.path.join(config_root, f"baseline_fd{fd}.json")
30 |         arch_config = building.load_config(arch_config_path)
31 |         pre_config_path = os.path.join(config_root, f"baseline_pre_fd{fd}.json")
32 |         pre_config = building.load_config(pre_config_path)
33 |         for num_fails, fails in enumerate(percent_fails):
34 |             seed = seeds[num_fd][num_fails]
35 |             try:
36 |                 run(
37 |                     fd,
38 |                     None,
39 |                     fails,
40 |                     arch_config,
41 |                     pre_config,
42 |                     pretrain=False,
43 |                     encoder=encoder,
44 |                     mode="metric",
45 |                     record_embeddings=False,
46 |                     replications=10,
47 |                     gpu=[0],
48 |                     master_seed=seed,
49 |                     version=f"{base_version}_baseline@{fails:.2f}",
50 |                 )
51 |             except Exception as e:
52 |                 error_log.append(e)
53 | 
54 |     ray.shutdown()
55 | 
56 |     if error_log:
57 |         for e in error_log:
58 |             print(type(e), e)
59 |     else:
60 |         print("Everything went well.")
61 | 
62 | 
63 | if __name__ == "__main__":
64 |     import argparse
65 | 
66 |     parser = argparse.ArgumentParser(
67 |         description="Reproduce all supervised baseline experiments."
68 |     )
69 |     parser.add_argument(
70 |         "base_version", default=None, help="common prefix for the version tag"
71 |     )
72 |     parser.add_argument(
73 |         "--percent_fails",
74 |         nargs="*",
75 |         default=[1.0, 0.4, 0.2, 0.1, 0.02],
76 |         help="percentage of failed runs",
77 |     )
78 |     parser.add_argument(
79 |         "--encoder",
80 |         default="cnn",
81 |         choices=["cnn", "lstm"],
82 |         help="encoder type",
83 |     )
84 |     parser.add_argument(
85 |         "--seed", default=42, help="master seed used to produce all other seeds"
86 |     )
87 |     opt = parser.parse_args()
88 | 
89 |     reproduce(opt.base_version, opt.percent_fails, opt.encoder, opt.seed)
90 | 


--------------------------------------------------------------------------------
/src/reproduce_semi_supervised.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | import random
  3 | from datetime import datetime
  4 | 
  5 | import ray
  6 | 
  7 | import building
  8 | from run_rbm import run as run_rbm
  9 | from run_semi_supervised import run as run_ae_or_metric
 10 | 
 11 | script_path = os.path.dirname(__file__)
 12 | config_root = os.path.join(script_path, "..", "configs")
 13 | 
 14 | 
 15 | def reproduce(base_version, percent_fails, percent_broken, encoder, mode, master_seed):
 16 |     ray.init()
 17 | 
 18 |     if base_version is None:
 19 |         base_version = datetime.now().timestamp()
 20 |     error_log = []
 21 |     fds = [1, 2, 3, 4]
 22 |     pre_mode = "pre" if mode == "metric" else "ae"
 23 |     random.seed(master_seed)
 24 |     seeds = [
 25 |         [
 26 |             [random.randint(0, 99999999) for _ in range(len(percent_fails))]
 27 |             for _ in range(len(percent_broken))
 28 |         ]
 29 |         for _ in range(len(fds))
 30 |     ]
 31 | 
 32 |     for num_fd, fd in enumerate(fds):
 33 |         arch_config_path = os.path.join(config_root, f"{encoder}_fd{fd}.json")
 34 |         arch_config = building.load_config(arch_config_path)
 35 |         pre_config_path = os.path.join(config_root, f"{encoder}_{pre_mode}_fd{fd}.json")
 36 |         pre_config = building.load_config(pre_config_path)
 37 |         for num_broken, broken in enumerate(percent_broken):
 38 |             for num_fails, fails in enumerate(percent_fails):
 39 |                 seed = seeds[num_fd][num_broken][num_fails]
 40 |                 try:
 41 |                     run(
 42 |                         fd,
 43 |                         broken,
 44 |                         fails,
 45 |                         arch_config,
 46 |                         pre_config,
 47 |                         encoder,
 48 |                         mode,
 49 |                         seed,
 50 |                         base_version,
 51 |                     )
 52 |                 except Exception as e:
 53 |                     error_log.append(e)
 54 | 
 55 |     ray.shutdown()
 56 | 
 57 |     if error_log:
 58 |         for e in error_log:
 59 |             print(type(e), e)
 60 |     else:
 61 |         print("Everything went well.")
 62 | 
 63 | 
 64 | def run(fd, broken, fails, arch_config, pre_config, encoder, mode, seed, base_version):
 65 |     if mode == "rbm":
 66 |         run_rbm(
 67 |             fd,
 68 |             broken,
 69 |             fails,
 70 |             arch_config,
 71 |             encoder,
 72 |             replications=10,
 73 |             gpu=0,
 74 |             master_seed=seed,
 75 |             version=f"{base_version}_pretrained@{broken:.2f}@{fails:.2f}",
 76 |         )
 77 |     else:
 78 |         run_ae_or_metric(
 79 |             fd,
 80 |             broken,
 81 |             fails,
 82 |             arch_config,
 83 |             pre_config,
 84 |             pretrain=True,
 85 |             encoder=encoder,
 86 |             mode=mode,
 87 |             record_embeddings=False,
 88 |             replications=10,
 89 |             gpu=[0],
 90 |             master_seed=seed,
 91 |             version=f"{base_version}_pretrained@{broken:.2f}@{fails:.2f}",
 92 |         )
 93 | 
 94 | 
 95 | if __name__ == "__main__":
 96 |     import argparse
 97 | 
 98 |     parser = argparse.ArgumentParser(
 99 |         description="Reproduce all semi-supervised experiments."
100 |     )
101 |     parser.add_argument(
102 |         "base_version", default=None, help="common prefix for the version tag"
103 |     )
104 |     parser.add_argument(
105 |         "--percent_fails",
106 |         nargs="*",
107 |         default=[1.0, 0.4, 0.2, 0.1, 0.02],
108 |         help="percentage of failed runs",
109 |     )
110 |     parser.add_argument(
111 |         "--percent_broken",
112 |         nargs="*",
113 |         type=float,
114 |         default=[0.9, 0.8, 0.7, 0.6, 0.4],
115 |         help="percentage of degradation of unfailed runs",
116 |     )
117 |     parser.add_argument(
118 |         "--encoder",
119 |         default="cnn",
120 |         choices=["cnn", "lstm"],
121 |         help="encoder type",
122 |     )
123 |     parser.add_argument(
124 |         "--mode",
125 |         default="metric",
126 |         choices=["metric", "autoencoder", "rbm"],
127 |         help="metric or autoencoder pre-training mode",
128 |     )
129 |     parser.add_argument(
130 |         "--seed", default=21, help="master seed used to produce all other seeds"
131 |     )
132 |     opt = parser.parse_args()
133 | 
134 |     reproduce(
135 |         opt.base_version,
136 |         opt.percent_fails,
137 |         opt.percent_broken,
138 |         opt.encoder,
139 |         opt.mode,
140 |         opt.seed,
141 |     )
142 | 


--------------------------------------------------------------------------------
/src/run_baseline.py:
--------------------------------------------------------------------------------
 1 | import random
 2 | from datetime import datetime
 3 | 
 4 | import building
 5 | 
 6 | 
 7 | def run(
 8 |     source,
 9 |     fails,
10 |     config,
11 |     encoder,
12 |     seed,
13 |     gpu,
14 |     pretrained_encoder_path,
15 |     version,
16 |     record_embeddings=False,
17 | ):
18 |     trainer, data, model = building.build_baseline(
19 |         source,
20 |         fails,
21 |         config,
22 |         encoder,
23 |         pretrained_encoder_path,
24 |         gpu,
25 |         seed,
26 |         version,
27 |         record_embeddings,
28 |     )
29 |     trainer.fit(model, datamodule=data)
30 |     trainer.test(datamodule=data)
31 | 
32 | 
33 | def run_multiple(
34 |     source,
35 |     fails,
36 |     config,
37 |     encoder,
38 |     replications,
39 |     gpu,
40 |     pretrained_encoder_path,
41 |     version=None,
42 | ):
43 |     random.seed(999)
44 |     seeds = [random.randint(0, 9999999) for _ in range(replications)]
45 |     if version is None:
46 |         version = datetime.now().timestamp()
47 | 
48 |     for f in fails:
49 |         for s in seeds:
50 |             run(source, f, config, encoder, s, gpu, pretrained_encoder_path, version)
51 | 
52 | 
53 | if __name__ == "__main__":
54 |     import argparse
55 | 
56 |     parser = argparse.ArgumentParser(description="Run baseline experiment")
57 |     parser.add_argument("--source", type=int, help="FD number of the source data")
58 |     parser.add_argument(
59 |         "-f", "--fails", nargs="+", type=float, help="percent fail runs to use"
60 |     )
61 |     parser.add_argument("--config", required=True, help="path to config file")
62 |     parser.add_argument(
63 |         "--encoder",
64 |         default="cnn",
65 |         choices=["cnn", "lstm"],
66 |         help="encoder type",
67 |     )
68 |     parser.add_argument(
69 |         "--pretrained_encoder",
70 |         default=None,
71 |         help="Path to checkpoint file form pretraining",
72 |     )
73 |     parser.add_argument(
74 |         "-r", "--replications", type=int, default=3, help="replications for each run"
75 |     )
76 |     parser.add_argument("--gpu", type=int, default=0, help="id of GPU to use")
77 |     opt = parser.parse_args()
78 | 
79 |     _config = building.load_config(opt.config)
80 |     run_multiple(
81 |         opt.source,
82 |         opt.fails,
83 |         _config,
84 |         opt.encoder,
85 |         opt.replications,
86 |         opt.gpu,
87 |         opt.pretrained_encoder,
88 |     )
89 | 


--------------------------------------------------------------------------------
/src/run_pretraining.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | import random
  3 | from datetime import datetime
  4 | 
  5 | import building
  6 | 
  7 | 
  8 | def run(
  9 |     source,
 10 |     target,
 11 |     percent_broken,
 12 |     percent_fail_runs,
 13 |     arch_config,
 14 |     config,
 15 |     encoder,
 16 |     mode,
 17 |     record_embeddings,
 18 |     seed,
 19 |     gpu,
 20 |     version,
 21 | ):
 22 |     trainer, data, model = building.build_pretraining(
 23 |         source,
 24 |         target,
 25 |         percent_broken,
 26 |         percent_fail_runs,
 27 |         arch_config,
 28 |         config,
 29 |         encoder,
 30 |         mode,
 31 |         record_embeddings,
 32 |         gpu,
 33 |         seed,
 34 |         version,
 35 |     )
 36 |     trainer.fit(model, datamodule=data)
 37 |     trainer.test(datamodule=data)
 38 | 
 39 |     checkpoint_path = _get_checkpoint_path(trainer.logger)
 40 |     best_score = trainer.checkpoint_callback.best_model_score
 41 | 
 42 |     return checkpoint_path, best_score
 43 | 
 44 | 
 45 | def _get_checkpoint_path(logger):
 46 |     checkpoints_path = logger.checkpoint_path
 47 |     *_, checkpoint = sorted(
 48 |         [f for f in os.listdir(checkpoints_path)]
 49 |     )  # get last checkpoint
 50 |     checkpoint_path = os.path.join(checkpoints_path, checkpoint)
 51 | 
 52 |     return checkpoint_path
 53 | 
 54 | 
 55 | def run_multiple(
 56 |     source,
 57 |     target,
 58 |     broken,
 59 |     fails,
 60 |     arch_config,
 61 |     config,
 62 |     encoder,
 63 |     mode,
 64 |     record_embeddings,
 65 |     replications,
 66 |     gpu,
 67 |     version=None,
 68 | ):
 69 |     broken = broken if broken is not None else [1.0]
 70 |     random.seed(999)
 71 |     seeds = [random.randint(0, 9999999) for _ in range(replications)]
 72 |     if version is None:
 73 |         version = datetime.now().timestamp()
 74 | 
 75 |     checkpoints = {b: [] for b in broken}
 76 |     for b in broken:
 77 |         for f in fails:
 78 |             for s in seeds:
 79 |                 checkpoint_path = run(
 80 |                     source,
 81 |                     target,
 82 |                     b,
 83 |                     f,
 84 |                     arch_config,
 85 |                     config,
 86 |                     encoder,
 87 |                     mode,
 88 |                     record_embeddings,
 89 |                     s,
 90 |                     gpu,
 91 |                     version,
 92 |                 )
 93 |                 checkpoints[b].append(checkpoint_path)
 94 | 
 95 |     return checkpoints
 96 | 
 97 | 
 98 | if __name__ == "__main__":
 99 |     import argparse
100 | 
101 |     parser = argparse.ArgumentParser(description="Run domain adaption experiment")
102 |     parser.add_argument(
103 |         "--source", type=int, required=True, help="FD number of the source data"
104 |     )
105 |     parser.add_argument("--target", type=int, help="FD number of the target data")
106 |     parser.add_argument(
107 |         "-b", "--broken", nargs="+", type=float, help="percent broken to use"
108 |     )
109 |     parser.add_argument(
110 |         "-f", "--fails", nargs="+", type=float, help="percent fail runs to use"
111 |     )
112 |     parser.add_argument(
113 |         "--arch_config",
114 |         required=True,
115 |         help="path to architecture config file",
116 |     )
117 |     parser.add_argument("--config", required=True, help="path to config file")
118 |     parser.add_argument(
119 |         "--encoder",
120 |         default="cnn",
121 |         choices=["cnn", "lstm"],
122 |         help="encoder type",
123 |     )
124 |     parser.add_argument(
125 |         "--mode",
126 |         default="metric",
127 |         choices=["metric", "autoencoder"],
128 |         help="metric or autoencoder pre-training mode",
129 |     )
130 |     parser.add_argument(
131 |         "--record_embeddings",
132 |         action="store_true",
133 |         help="whether to record embeddings of val data",
134 |     )
135 |     parser.add_argument(
136 |         "-r", "--replications", type=int, default=3, help="replications for each run"
137 |     )
138 |     parser.add_argument("--gpu", type=int, default=0, help="id of GPU to use")
139 |     opt = parser.parse_args()
140 | 
141 |     _arch_config = building.load_config(opt.arch_config)
142 |     _config = building.load_config(opt.config)
143 |     run_multiple(
144 |         opt.source,
145 |         opt.target,
146 |         opt.broken,
147 |         opt.fails,
148 |         _arch_config,
149 |         _config,
150 |         opt.encoder,
151 |         opt.mode,
152 |         opt.record_embeddings,
153 |         opt.replications,
154 |         opt.gpu,
155 |     )
156 | 


--------------------------------------------------------------------------------
/src/run_rbm.py:
--------------------------------------------------------------------------------
  1 | import random
  2 | from datetime import datetime
  3 | 
  4 | import ray
  5 | import torch.optim
  6 | from pytorch_probgraph import RestrictedBoltzmannMachineCD
  7 | from sklearn.model_selection import ShuffleSplit
  8 | 
  9 | import building
 10 | from datasets.loader import CMAPSSLoader
 11 | 
 12 | 
 13 | def run(
 14 |     source,
 15 |     percent_broken,
 16 |     percent_fails,
 17 |     arch_config,
 18 |     encoder,
 19 |     replications,
 20 |     gpu,
 21 |     master_seed,
 22 |     version=None,
 23 | ):
 24 |     if version is None:
 25 |         version = datetime.now().timestamp()
 26 |     if master_seed:
 27 |         random.seed(master_seed)
 28 |     seeds = [random.randint(0, 9999999) for _ in range(replications)]
 29 | 
 30 |     if percent_fails < 1:
 31 |         splitter = ShuffleSplit(
 32 |             n_splits=replications, train_size=percent_fails, random_state=42
 33 |         )
 34 |     else:
 35 |         splitter = AllDataSplitter(replications)
 36 |     run_idx = range(CMAPSSLoader.NUM_TRAIN_RUNS[source])
 37 |     process_ids = []
 38 |     for (failed_idx, _), s in zip(splitter.split(run_idx), seeds):
 39 |         process_ids.append(
 40 |             ray_train.remote(
 41 |                 source,
 42 |                 percent_broken,
 43 |                 failed_idx,
 44 |                 arch_config,
 45 |                 encoder,
 46 |                 s,
 47 |                 gpu,
 48 |                 version,
 49 |             )
 50 |         )
 51 | 
 52 |     ray.get(process_ids)
 53 | 
 54 | 
 55 | class AllDataSplitter:
 56 |     """Splitter that returns the whole dataset as training and nothing as test split."""
 57 | 
 58 |     def __init__(self, n_splits):
 59 |         self.n_splits = n_splits
 60 | 
 61 |     def split(self, x):
 62 |         for _ in range(self.n_splits):
 63 |             yield list(x), []
 64 | 
 65 | 
 66 | @ray.remote(num_cpus=3, num_gpus=0.5)
 67 | def ray_train(source, percent_broken, failed_idx, arch_config, encoder, s, gpu, version):
 68 |     pretrained_rbm = run_pretraining(
 69 |         source,
 70 |         percent_broken,
 71 |         failed_idx,
 72 |         arch_config,
 73 |         gpu,
 74 |     )
 75 |     run_baseline(
 76 |         source,
 77 |         failed_idx,
 78 |         arch_config,
 79 |         encoder,
 80 |         s,
 81 |         gpu,
 82 |         pretrained_rbm,
 83 |         version,
 84 |     )
 85 | 
 86 | 
 87 | def run_pretraining(source, percent_broken, failed_idx, arch_config, gpu):
 88 |     print("Pre-train RBM layer...")
 89 |     device = torch.device(f"cuda:{gpu}")
 90 |     dm = building.build_datamodule(
 91 |         source,
 92 |         None,
 93 |         percent_broken,
 94 |         failed_idx,
 95 |         arch_config["batch_size"],
 96 |         truncate_val=True,
 97 |         distance_mode="linear",
 98 |         min_distance=1,
 99 |     )
100 |     dm.prepare_data()
101 |     dm.setup()
102 |     data = (torch.cat([anchor, query]) for anchor, query, *_ in dm.train_dataloader())
103 |     rbm: RestrictedBoltzmannMachineCD = building.build_rbm(
104 |         14, arch_config["base_filters"]
105 |     ).to(device)
106 |     adam = torch.optim.Adam(rbm.parameters(), lr=1e-4)
107 |     rbm.train(data, epochs=5, optimizer=adam, device=device)
108 | 
109 |     data = (torch.cat([anchor, query]) for anchor, query, *_ in dm.val_dataloader()[0])
110 |     recon_error = 0
111 |     num_elem = 0
112 |     with torch.no_grad():
113 |         for batch in data:
114 |             batch = batch.to(device)
115 |             recon = rbm.reconstruct(visible_input=batch)
116 |             recon_error += torch.sum((batch - recon) ** 2).cpu().item()
117 |             num_elem += len(batch)
118 |         recon_error /= num_elem * dm.window_size
119 | 
120 |     print(f"Recon error: {recon_error:.2f}")
121 | 
122 |     return rbm
123 | 
124 | 
125 | def run_baseline(source, fails, config, encoder, seed, gpu, pretrained_rbm, version):
126 |     trainer, data, model = building.build_baseline(
127 |         source, fails, config, encoder, None, gpu, seed, version
128 |     )
129 |     model.hparams["encoder"] = f"rbm_{encoder}"
130 |     model.load_from_rbm(pretrained_rbm)
131 |     trainer.fit(model, datamodule=data)
132 |     trainer.test(datamodule=data)
133 | 
134 | 
135 | if __name__ == "__main__":
136 |     import argparse
137 | 
138 |     parser = argparse.ArgumentParser(description="Run semi-supervised experiment")
139 |     parser.add_argument(
140 |         "--source", required=True, type=int, help="FD number of the source data"
141 |     )
142 |     parser.add_argument(
143 |         "-b", "--broken", required=True, type=float, help="percent broken to use"
144 |     )
145 |     parser.add_argument(
146 |         "-f", "--fails", required=True, type=float, help="percent fail runs to use"
147 |     )
148 |     parser.add_argument(
149 |         "--arch_config", required=True, help="path to architecture config file"
150 |     )
151 |     parser.add_argument(
152 |         "-r",
153 |         "--replications",
154 |         type=int,
155 |         default=1,
156 |         help="runs of the cross-validation",
157 |     )
158 |     parser.add_argument(
159 |         "--encoder",
160 |         default="cnn",
161 |         choices=["cnn", "lstm"],
162 |         help="encoder type",
163 |     )
164 |     parser.add_argument("--gpu", type=int, default=0, help="id of GPU to use")
165 |     parser.add_argument(
166 |         "--seed", default=999, help="master seed used to produce all seeds"
167 |     )
168 |     parser.add_argument("--version", help="version tag to group runs together")
169 |     opt = parser.parse_args()
170 | 
171 |     _arch_config = building.load_config(opt.arch_config)
172 |     ray.init()
173 |     run(
174 |         opt.source,
175 |         opt.broken,
176 |         opt.fails,
177 |         _arch_config,
178 |         opt.encoder,
179 |         opt.replications,
180 |         opt.gpu,
181 |         opt.seed,
182 |         opt.version,
183 |     )
184 |     ray.shutdown()
185 | 


--------------------------------------------------------------------------------
/src/run_semi_supervised.py:
--------------------------------------------------------------------------------
  1 | import random
  2 | from datetime import datetime
  3 | 
  4 | import ray
  5 | from sklearn.model_selection import ShuffleSplit
  6 | 
  7 | import building
  8 | from datasets.loader import CMAPSSLoader
  9 | from run_baseline import run as run_baseline
 10 | from run_pretraining import run as run_pretraining
 11 | 
 12 | 
 13 | def run(
 14 |     source,
 15 |     percent_broken,
 16 |     percent_fails,
 17 |     arch_config,
 18 |     pre_config,
 19 |     pretrain,
 20 |     encoder,
 21 |     mode,
 22 |     record_embeddings,
 23 |     replications,
 24 |     gpu,
 25 |     master_seed,
 26 |     version=None,
 27 | ):
 28 |     if version is None:
 29 |         version = datetime.now().timestamp()
 30 |     if master_seed:
 31 |         random.seed(master_seed)
 32 |     seeds = [random.randint(0, 9999999) for _ in range(replications)]
 33 | 
 34 |     if percent_fails < 1:
 35 |         splitter = ShuffleSplit(
 36 |             n_splits=replications, train_size=percent_fails, random_state=42
 37 |         )
 38 |     else:
 39 |         splitter = AllDataSplitter(replications)
 40 |     run_idx = range(CMAPSSLoader.NUM_TRAIN_RUNS[source])
 41 |     process_ids = []
 42 |     for (failed_idx, _), s in zip(splitter.split(run_idx), seeds):
 43 |         process_ids.append(
 44 |             ray_train.remote(
 45 |                 source,
 46 |                 percent_broken,
 47 |                 failed_idx,
 48 |                 arch_config,
 49 |                 pre_config,
 50 |                 encoder,
 51 |                 mode,
 52 |                 record_embeddings,
 53 |                 s,
 54 |                 gpu,
 55 |                 version,
 56 |                 pretrain,
 57 |             )
 58 |         )
 59 | 
 60 |     ray.get(process_ids)
 61 | 
 62 | 
 63 | class AllDataSplitter:
 64 |     """Splitter that returns the whole dataset as training and nothing as test split."""
 65 | 
 66 |     def __init__(self, n_splits):
 67 |         self.n_splits = n_splits
 68 | 
 69 |     def split(self, x):
 70 |         for _ in range(self.n_splits):
 71 |             yield list(x), []
 72 | 
 73 | 
 74 | @ray.remote(num_cpus=3, num_gpus=0.5)
 75 | def ray_train(
 76 |     source,
 77 |     percent_broken,
 78 |     failed_idx,
 79 |     arch_config,
 80 |     pre_config,
 81 |     encoder,
 82 |     mode,
 83 |     record_embeddings,
 84 |     s,
 85 |     gpu,
 86 |     version,
 87 |     pretrain,
 88 | ):
 89 |     if pretrain:
 90 |         checkpoint, _ = run_pretraining(
 91 |             source,
 92 |             None,
 93 |             percent_broken,
 94 |             failed_idx,
 95 |             arch_config,
 96 |             pre_config,
 97 |             encoder,
 98 |             mode,
 99 |             record_embeddings,
100 |             s,
101 |             gpu,
102 |             version,
103 |         )
104 |     else:
105 |         checkpoint = None
106 |     run_baseline(
107 |         source,
108 |         failed_idx,
109 |         arch_config,
110 |         encoder,
111 |         s,
112 |         gpu,
113 |         checkpoint,
114 |         version,
115 |         record_embeddings,
116 |     )
117 | 
118 | 
119 | if __name__ == "__main__":
120 |     import argparse
121 | 
122 |     parser = argparse.ArgumentParser(description="Run semi-supervised experiment")
123 |     parser.add_argument(
124 |         "--source", required=True, type=int, help="FD number of the source data"
125 |     )
126 |     parser.add_argument(
127 |         "-b", "--broken", required=True, type=float, help="percent broken to use"
128 |     )
129 |     parser.add_argument(
130 |         "-f", "--fails", required=True, type=float, help="percent fail runs to use"
131 |     )
132 |     parser.add_argument(
133 |         "--arch_config", required=True, help="path to architecture config file"
134 |     )
135 |     parser.add_argument(
136 |         "--pre_config", required=True, help="path to pretraining config file"
137 |     )
138 |     parser.add_argument("--pretrain", action="store_true", help="use pre-training")
139 |     parser.add_argument(
140 |         "-r",
141 |         "--replications",
142 |         type=int,
143 |         default=1,
144 |         help="runs of the cross-validation",
145 |     )
146 |     parser.add_argument(
147 |         "--encoder",
148 |         default="cnn",
149 |         choices=["cnn", "lstm"],
150 |         help="encoder type",
151 |     )
152 |     parser.add_argument(
153 |         "--mode",
154 |         default="metric",
155 |         choices=["metric", "autoencoder"],
156 |         help="metric or autoencoder pre-training mode",
157 |     )
158 |     parser.add_argument(
159 |         "--record_embeddings",
160 |         action="store_true",
161 |         help="whether to record embeddings of val data",
162 |     )
163 |     parser.add_argument("--gpu", type=int, default=0, help="id of GPU to use")
164 |     parser.add_argument(
165 |         "--seed", default=999, help="master seed used to produce all seeds"
166 |     )
167 |     parser.add_argument("--version", help="version tag to group runs together")
168 |     opt = parser.parse_args()
169 | 
170 |     _arch_config = building.load_config(opt.arch_config)
171 |     _pre_config = building.load_config(opt.pre_config)
172 |     ray.init()
173 |     run(
174 |         opt.source,
175 |         opt.broken,
176 |         opt.fails,
177 |         _arch_config,
178 |         _pre_config,
179 |         opt.pretrain,
180 |         opt.encoder,
181 |         opt.mode,
182 |         opt.record_embeddings,
183 |         opt.replications,
184 |         opt.gpu,
185 |         opt.seed,
186 |         opt.version,
187 |     )
188 |     ray.shutdown()
189 | 


--------------------------------------------------------------------------------
/tests/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/tilman151/self-supervised-ssl/745d4b8b06a946f1be1e2c8bc15e4b0b55c7c4ff/tests/__init__.py


--------------------------------------------------------------------------------
/tests/dataset_tests/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/tilman151/self-supervised-ssl/745d4b8b06a946f1be1e2c8bc15e4b0b55c7c4ff/tests/dataset_tests/__init__.py


--------------------------------------------------------------------------------
/tests/dataset_tests/templates.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | from torch.utils.data import TensorDataset
 3 | 
 4 | from datasets.cmapss import PairedCMAPSS
 5 | 
 6 | 
 7 | class PretrainingDataModuleTemplate:
 8 |     def test_data_structure(self):
 9 |         with self.subTest(split="dev"):
10 |             dataloader = self.dataset.train_dataloader()
11 |             self._check_paired_dataset(dataloader.dataset)
12 | 
13 |         with self.subTest(split="val"):
14 |             loaders = self.dataset.val_dataloader()
15 |             self.assertIsInstance(loaders, list)
16 |             self.assertEqual(self.expected_num_val_loaders, len(loaders))
17 |             self._check_paired_dataset(loaders[0].dataset)
18 |             for dataloader in loaders[1:]:
19 |                 self._check_tensor_dataset(dataloader.dataset)
20 | 
21 |     def _check_paired_dataset(self, data):
22 |         self.assertIsInstance(data, PairedCMAPSS)
23 |         self._check_paired_shapes(data)
24 | 
25 |     def _check_paired_shapes(self, data):
26 |         for anchors, queries, distances, domain_labels in data:
27 |             self.assertEqual(torch.Size((14, self.window_size)), anchors.shape)
28 |             self.assertEqual(torch.Size((14, self.window_size)), queries.shape)
29 |             self.assertEqual(torch.Size(()), distances.shape)
30 |             self.assertEqual(torch.Size(()), domain_labels.shape)
31 | 
32 |     def _check_tensor_dataset(self, data):
33 |         self.assertIsInstance(data, TensorDataset)
34 |         self._check_cmapss_shapes(data)
35 | 
36 |     def _check_cmapss_shapes(self, data):
37 |         for i in range(len(data)):
38 |             features, labels = data[i]
39 |             self.assertEqual(torch.Size((14, self.window_size)), features.shape)
40 |             self.assertEqual(torch.Size(()), labels.shape)
41 | 
42 |     def test_distances(self):
43 |         with self.subTest(split="dev"):
44 |             _, _, distances, _ = self._run_epoch(self.dataset.train_dataloader())
45 |             self.assertTrue(torch.all(distances >= 0))
46 | 
47 |         with self.subTest(split="val"):
48 |             _, _, distances, _ = self._run_epoch(self.dataset.val_dataloader()[0])
49 |             self.assertTrue(torch.all(distances >= 0))
50 | 
51 |     def test_determinism(self):
52 |         with self.subTest(split="dev"):
53 |             train_loader = self.dataset.train_dataloader()
54 |             *one_train_data, one_domain_labels = self._run_epoch(train_loader)
55 |             *another_train_data, another_domain_labels = self._run_epoch(train_loader)
56 | 
57 |             for one, another in zip(one_train_data, another_train_data):
58 |                 self.assertNotEqual(0.0, torch.sum(one - another))
59 | 
60 |         with self.subTest(split="val"):
61 |             paired_val_loader = self.dataset.val_dataloader()[0]
62 |             one_train_data = self._run_epoch(paired_val_loader)
63 |             another_train_data = self._run_epoch(paired_val_loader)
64 | 
65 |             for one, another in zip(one_train_data, another_train_data):
66 |                 self.assertEqual(0.0, torch.sum(one - another))
67 | 
68 |     def _run_epoch(self, loader):
69 |         anchors = torch.empty((len(loader.dataset), 14, self.window_size))
70 |         queries = torch.empty((len(loader.dataset), 14, self.window_size))
71 |         distances = torch.empty(len(loader.dataset))
72 |         domain_labels = torch.empty(len(loader.dataset))
73 | 
74 |         start = 0
75 |         end = loader.batch_size
76 |         for anchor, query, dist, domain in loader:
77 |             anchors[start:end] = anchor
78 |             queries[start:end] = query
79 |             distances[start:end] = dist
80 |             domain_labels[start:end] = domain
81 |             start = end
82 |             end += anchor.shape[0]
83 | 
84 |         return anchors, queries, distances, domain_labels
85 | 
86 |     def test_min_distance(self):
87 |         train_loader = self.dataset.train_dataloader()
88 |         self.assertEqual(self.dataset.min_distance, train_loader.dataset.min_distance)
89 |         val_loader = self.dataset.val_dataloader()[0]
90 |         self.assertEqual(1, val_loader.dataset.min_distance)
91 | 


--------------------------------------------------------------------------------
/tests/dataset_tests/test_baseline.py:
--------------------------------------------------------------------------------
  1 | import unittest
  2 | 
  3 | import torch
  4 | import torch.utils.data
  5 | from torch.utils.data import RandomSampler, SequentialSampler
  6 | 
  7 | import datasets
  8 | from tests.dataset_tests.templates import PretrainingDataModuleTemplate
  9 | 
 10 | 
 11 | class TestCMAPSSBaseline(unittest.TestCase):
 12 |     def setUp(self):
 13 |         self.dataset = datasets.BaselineDataModule(
 14 |             3, batch_size=16, percent_fail_runs=0.8
 15 |         )
 16 |         self.dataset.prepare_data()
 17 |         self.dataset.setup()
 18 | 
 19 |     def test_override_window_size(self):
 20 |         dataset = datasets.BaselineDataModule(
 21 |             3, batch_size=16, window_size=40, percent_fail_runs=0.8
 22 |         )
 23 |         for fd in dataset.cmapss.values():
 24 |             self.assertEqual(40, fd.window_size)
 25 | 
 26 |     def test_default_window_size(self):
 27 |         window_sizes = [30, 20, 30, 15]
 28 |         for i, win in enumerate(window_sizes, start=1):
 29 |             dataset = datasets.BaselineDataModule(i, batch_size=16, percent_fail_runs=0.8)
 30 |             for fd in dataset.cmapss.values():
 31 |                 self.assertEqual(win, fd.window_size)
 32 | 
 33 |     def test_train_batch_structure(self):
 34 |         train_loader = self.dataset.train_dataloader()
 35 |         self.assertIsInstance(train_loader.sampler, RandomSampler)
 36 |         self._assert_batch_structure(train_loader)
 37 | 
 38 |     def test_val_batch_structure(self):
 39 |         val_loader = self.dataset.val_dataloader()
 40 |         self.assertIsInstance(val_loader.sampler, SequentialSampler)
 41 |         self._assert_batch_structure(val_loader)
 42 | 
 43 |     def test_test_batch_structure(self):
 44 |         test_loaders = self.dataset.test_dataloader()
 45 |         for test_loader in test_loaders:
 46 |             self.assertIsInstance(test_loader.sampler, SequentialSampler)
 47 |             self._assert_batch_structure(test_loader)
 48 | 
 49 |     def _assert_batch_structure(self, loader):
 50 |         batch = next(iter(loader))
 51 |         self.assertEqual(2, len(batch))
 52 |         features, labels = batch
 53 |         self.assertEqual(torch.Size((16, 14, 30)), features.shape)
 54 |         self.assertEqual(torch.Size((16,)), labels.shape)
 55 | 
 56 |     def test_selected_source_on_train(self):
 57 |         fd_source = self.dataset.fd_source
 58 |         baseline_train_dataset = self.dataset.train_dataloader().dataset
 59 |         source_train_dataset = self.dataset.cmapss[fd_source].train_dataloader().dataset
 60 |         self._assert_datasets_equal(baseline_train_dataset, source_train_dataset)
 61 | 
 62 |     def test_selected_source_on_val(self):
 63 |         fd_source = self.dataset.fd_source
 64 |         baseline_val_dataset = self.dataset.val_dataloader().dataset
 65 |         source_val_dataset = self.dataset.cmapss[fd_source].val_dataloader().dataset
 66 |         self._assert_datasets_equal(baseline_val_dataset, source_val_dataset)
 67 | 
 68 |     def test_selected_all_on_test(self):
 69 |         baseline_test_loaders = self.dataset.test_dataloader()
 70 |         for fd, baseline_test_loader in enumerate(baseline_test_loaders, start=1):
 71 |             baseline_test_dataset = baseline_test_loader.dataset
 72 |             test_dataset = self.dataset.cmapss[fd].test_dataloader().dataset
 73 |             self._assert_datasets_equal(baseline_test_dataset, test_dataset)
 74 | 
 75 |     def _assert_datasets_equal(self, baseline_dataset, inner_dataset):
 76 |         num_samples = len(baseline_dataset)
 77 |         baseline_data = baseline_dataset[:num_samples]
 78 |         inner_data = inner_dataset[:num_samples]
 79 |         for baseline, inner in zip(baseline_data, inner_data):
 80 |             self.assertEqual(0, torch.sum(baseline - inner))
 81 | 
 82 |     def test_fail_runs_passed_correctly(self):
 83 |         for i in range(1, 5):
 84 |             if i == self.dataset.fd_source:
 85 |                 self.assertEqual(0.8, self.dataset.cmapss[i].percent_fail_runs)
 86 |             else:
 87 |                 self.assertIsNone(self.dataset.cmapss[i].percent_fail_runs)
 88 | 
 89 | 
 90 | class TestPretrainingBaselineDataModuleFullData(
 91 |     unittest.TestCase, PretrainingDataModuleTemplate
 92 | ):
 93 |     def setUp(self):
 94 |         self.dataset = datasets.PretrainingBaselineDataModule(
 95 |             3, num_samples=10000, batch_size=16, min_distance=2
 96 |         )
 97 |         self.dataset.prepare_data()
 98 |         self.dataset.setup()
 99 | 
100 |         self.expected_num_val_loaders = 2
101 |         self.window_size = self.dataset.broken_source_loader.window_size
102 | 
103 |     def test_val_truncation(self):
104 |         with self.subTest(truncation=False):
105 |             dataset = datasets.PretrainingBaselineDataModule(
106 |                 3, num_samples=10000, batch_size=16
107 |             )
108 |             self.assertFalse(dataset.broken_source_loader.truncate_val)
109 |             self.assertFalse(dataset.source.truncate_val)
110 | 
111 |         with self.subTest(truncation=True):
112 |             dataset = datasets.PretrainingBaselineDataModule(
113 |                 3, num_samples=10000, batch_size=16, truncate_val=True
114 |             )
115 |             self.assertTrue(dataset.broken_source_loader.truncate_val)
116 |             self.assertTrue(dataset.source.truncate_val)
117 | 
118 |     def test_override_window_size(self):
119 |         dataset = datasets.PretrainingBaselineDataModule(
120 |             3, num_samples=10000, batch_size=16, window_size=40
121 |         )
122 |         dataset.prepare_data()
123 |         dataset.setup()
124 |         train_loader = dataset.train_dataloader()
125 | 
126 |         anchors, queries, _, _ = next(iter(train_loader))
127 |         self.assertEqual(40, anchors.shape[2])
128 |         self.assertEqual(40, queries.shape[2])
129 | 
130 |     def test_truncation_passed_correctly(self):
131 |         dataset = datasets.PretrainingBaselineDataModule(
132 |             3, 1000, 16, percent_broken=0.2, percent_fail_runs=0.5
133 |         )
134 |         self.assertEqual(0.2, dataset.broken_source_loader.percent_broken)
135 |         self.assertIsNone(dataset.broken_source_loader.percent_fail_runs)
136 |         self.assertEqual(0.5, dataset.fails_source_loader.percent_fail_runs)
137 |         self.assertIsNone(dataset.fails_source_loader.percent_broken)
138 | 
139 |     def test_distance_mode_passed_correctly(self):
140 |         dataset = datasets.PretrainingBaselineDataModule(
141 |             3, 1000, 16, distance_mode="labeled"
142 |         )
143 |         data_loader = dataset.train_dataloader()
144 |         self.assertEqual(dataset.distance_mode, data_loader.dataset.mode)
145 | 
146 |     def test_both_source_datasets_used(self):
147 |         dataset = datasets.PretrainingBaselineDataModule(
148 |             3, 1000, 16, percent_broken=0.2, percent_fail_runs=0.5
149 |         )
150 |         for split in ["dev", "val"]:
151 |             with self.subTest(split):
152 |                 num_broken_runs = len(dataset.broken_source_loader.load_split(split)[0])
153 |                 num_fail_runs = len(dataset.fails_source_loader.load_split(split)[0])
154 |                 paired_dataset = dataset._get_paired_dataset(split)
155 |                 self.assertEqual(
156 |                     num_broken_runs + num_fail_runs, len(paired_dataset._features)
157 |                 )
158 | 
159 |     def test_get_unfailed_runs(self):
160 |         fail_runs = [1, 5, 40, 79]
161 |         dataset = datasets.PretrainingBaselineDataModule(
162 |             3, 1000, 16, percent_broken=0.2, percent_fail_runs=fail_runs
163 |         )
164 |         self.assertListEqual(fail_runs, dataset.fails_source_loader.percent_fail_runs)
165 |         for r in fail_runs:
166 |             self.assertNotIn(r, dataset.broken_source_loader.percent_fail_runs)
167 | 
168 |     def test_percent_fails_zero(self):
169 |         dataset = datasets.PretrainingBaselineDataModule(
170 |             3, 1000, 16, percent_broken=0.2, percent_fail_runs=0.0
171 |         )
172 |         num_broken_runs = len(dataset.broken_source_loader.load_split("dev")[0])
173 |         num_fail_runs = len(dataset.fails_source_loader.load_split("dev")[0])
174 |         self.assertEqual(0, num_fail_runs)
175 |         self.assertNotEqual(0, num_broken_runs)
176 | 
177 | 
178 | class TestPretrainingBaselineDataModuleLowData(
179 |     unittest.TestCase, PretrainingDataModuleTemplate
180 | ):
181 |     def setUp(self):
182 |         self.dataset = datasets.PretrainingBaselineDataModule(
183 |             3, num_samples=10000, batch_size=16, percent_broken=0.2
184 |         )
185 |         self.dataset.prepare_data()
186 |         self.dataset.setup()
187 | 
188 |         self.expected_num_val_loaders = 2
189 |         self.window_size = self.dataset.broken_source_loader.window_size
190 | 


--------------------------------------------------------------------------------
/tests/dataset_tests/test_cmapss.py:
--------------------------------------------------------------------------------
  1 | import unittest
  2 | from unittest import mock
  3 | 
  4 | import numpy as np
  5 | import torch
  6 | from torch.utils.data import TensorDataset
  7 | 
  8 | import datasets
  9 | from datasets import cmapss
 10 | from datasets import loader
 11 | 
 12 | 
 13 | class TestCMAPSS(unittest.TestCase):
 14 |     def test_data(self):
 15 |         window_sizes = [30, 20, 30, 15]
 16 |         for n, win in enumerate(window_sizes, start=1):
 17 |             dataset = cmapss.CMAPSSDataModule(n, batch_size=16)
 18 |             dataset.prepare_data()
 19 |             dataset.setup()
 20 |             for split in ["dev", "val", "test"]:
 21 |                 with self.subTest(fd=n, split=split):
 22 |                     features, targets = dataset.data[split]
 23 |                     self.assertEqual(win, features.shape[2])
 24 |                     self.assertEqual(len(features), len(targets))
 25 |                     self.assertEqual(torch.float32, features.dtype)
 26 |                     self.assertEqual(torch.float32, targets.dtype)
 27 | 
 28 |     def test_override_window_size(self):
 29 |         window_size = 40
 30 |         for n in range(1, 5):
 31 |             dataset = cmapss.CMAPSSDataModule(n, batch_size=16, window_size=window_size)
 32 |             self.assertEqual(window_size, dataset.window_size)
 33 |             self.assertEqual(window_size, dataset._loader.window_size)
 34 | 
 35 |     def test_default_window_size(self):
 36 |         window_sizes = [30, 20, 30, 15]
 37 |         for n, win in enumerate(window_sizes, start=1):
 38 |             dataset = cmapss.CMAPSSDataModule(n, batch_size=16)
 39 |             self.assertEqual(win, dataset.window_size)
 40 |             self.assertEqual(win, dataset._loader.window_size)
 41 | 
 42 |     def test_feature_select(self):
 43 |         feature_idx = [4, 9, 10, 13, 14, 15, 22]
 44 |         dataset = cmapss.CMAPSSDataModule(1, batch_size=16, feature_select=feature_idx)
 45 |         self.assertListEqual(feature_idx, dataset.feature_select)
 46 |         self.assertListEqual(feature_idx, dataset._loader.feature_select)
 47 | 
 48 |     def test_default_feature_select(self):
 49 |         dataset = cmapss.CMAPSSDataModule(1, batch_size=16)
 50 |         self.assertListEqual(dataset._loader.DEFAULT_CHANNELS, dataset.feature_select)
 51 |         self.assertListEqual(
 52 |             dataset._loader.DEFAULT_CHANNELS, dataset._loader.feature_select
 53 |         )
 54 | 
 55 |     def test_truncation_functions(self):
 56 |         full_dataset = cmapss.CMAPSSDataModule(fd=1, batch_size=4, window_size=30)
 57 |         full_dataset.prepare_data()
 58 |         full_dataset.setup()
 59 | 
 60 |         dataset = cmapss.CMAPSSDataModule(
 61 |             fd=1, batch_size=4, window_size=30, percent_fail_runs=0.8
 62 |         )
 63 |         dataset.prepare_data()
 64 |         dataset.setup()
 65 |         self.assertGreater(len(full_dataset.data["dev"][0]), len(dataset.data["dev"][0]))
 66 |         self.assertEqual(len(full_dataset.data["val"][0]), len(dataset.data["val"][0]))
 67 |         self.assertEqual(len(full_dataset.data["test"][0]), len(dataset.data["test"][0]))
 68 | 
 69 |         dataset = cmapss.CMAPSSDataModule(
 70 |             fd=1, batch_size=4, window_size=30, percent_broken=0.2
 71 |         )
 72 |         dataset.prepare_data()
 73 |         dataset.setup()
 74 |         self.assertGreater(len(full_dataset.data["dev"][0]), len(dataset.data["dev"][0]))
 75 |         self.assertAlmostEqual(
 76 |             0.2,
 77 |             len(dataset.data["dev"][0]) / len(full_dataset.data["dev"][0]),
 78 |             delta=0.01,
 79 |         )
 80 |         self.assertEqual(
 81 |             len(full_dataset.data["val"][0]), len(dataset.data["val"][0])
 82 |         )  # Val dataset_tests not truncated
 83 |         self.assertEqual(
 84 |             len(full_dataset.data["test"][0]), len(dataset.data["test"][0])
 85 |         )  # Test dataset_tests not truncated
 86 |         self.assertFalse(
 87 |             torch.any(dataset.data["dev"][1] == 1)
 88 |         )  # No failure dataset_tests in truncated dataset_tests
 89 |         self.assertEqual(
 90 |             full_dataset.data["dev"][1][0], dataset.data["dev"][1][0]
 91 |         )  # First target has to be equal
 92 | 
 93 |     def test_truncation_passed_correctly(self):
 94 |         dataset = cmapss.CMAPSSDataModule(1, 4, percent_broken=0.2, percent_fail_runs=0.5)
 95 |         self.assertEqual(dataset.percent_broken, dataset._loader.percent_broken)
 96 |         self.assertEqual(dataset.percent_fail_runs, dataset._loader.percent_fail_runs)
 97 | 
 98 |     def test_from_loader(self):
 99 |         cmapss_loader = loader.CMAPSSLoader(3, 40, 130, 0.2, 0.5, truncate_val=True)
100 |         cmapss_dataset = cmapss.CMAPSSDataModule.from_loader(cmapss_loader, 128)
101 |         self.assertEqual(cmapss_loader.fd, cmapss_dataset.fd)
102 |         self.assertEqual(cmapss_loader.window_size, cmapss_dataset.window_size)
103 |         self.assertEqual(cmapss_loader.max_rul, cmapss_dataset.max_rul)
104 |         self.assertEqual(cmapss_loader.percent_broken, cmapss_dataset.percent_broken)
105 |         self.assertEqual(
106 |             cmapss_loader.percent_fail_runs, cmapss_dataset.percent_fail_runs
107 |         )
108 |         self.assertEqual(cmapss_loader.feature_select, cmapss_dataset.feature_select)
109 |         self.assertEqual(cmapss_loader.truncate_val, cmapss_dataset.truncate_val)
110 |         self.assertEqual(128, cmapss_dataset.batch_size)
111 | 
112 |     def test_empty_dataset(self):
113 |         dataset = cmapss.CMAPSSDataModule(1, 4, percent_broken=0.2, percent_fail_runs=0.0)
114 |         dataset.setup()
115 |         dataset = cmapss.CMAPSSDataModule(1, 4, percent_broken=0.0, percent_fail_runs=0.5)
116 |         dataset.setup()
117 | 
118 | 
119 | class DummyCMAPSS:
120 |     def __init__(self, length):
121 |         self.window_size = 30
122 |         self.max_rul = 125
123 |         self.data = {
124 |             "dev": (
125 |                 [torch.zeros(length, self.window_size, 5)],
126 |                 [torch.clamp_max(torch.arange(length, 0, step=-1), 125)],
127 |             ),
128 |         }
129 | 
130 |     def load_split(self, split):
131 |         assert split == "dev", "Can only use dev data."
132 |         return self.data["dev"]
133 | 
134 | 
135 | class DummyCMAPSSShortRuns:
136 |     """Contains runs that are too short with zero features to distinguish them."""
137 | 
138 |     def __init__(self):
139 |         self.window_size = 30
140 |         self.max_rul = 125
141 |         self.data = {
142 |             "dev": (
143 |                 [
144 |                     torch.ones(100, self.window_size, 5)
145 |                     * torch.arange(1, 101).view(100, 1, 1),  # normal run
146 |                     torch.zeros(2, self.window_size, 5),  # too short run
147 |                     torch.ones(100, self.window_size, 5)
148 |                     * torch.arange(1, 101).view(100, 1, 1),  # normal run
149 |                     torch.zeros(1, self.window_size, 5),  # empty run
150 |                 ],
151 |                 [
152 |                     torch.clamp_max(torch.arange(100, 0, step=-1), 125),
153 |                     torch.ones(2) * 500,
154 |                     torch.clamp_max(torch.arange(100, 0, step=-1), 125),
155 |                     torch.ones(1) * 500,
156 |                 ],
157 |             ),
158 |         }
159 | 
160 |     def load_split(self, split):
161 |         assert split == "dev", "Can only use dev data."
162 |         return self.data["dev"]
163 | 
164 | 
165 | class TestPairedDataset(unittest.TestCase):
166 |     def setUp(self):
167 |         self.length = 300
168 |         self.cmapss_normal = DummyCMAPSS(self.length)
169 |         self.cmapss_short = DummyCMAPSSShortRuns()
170 |         self.fd1 = loader.CMAPSSLoader(1)
171 |         self.fd3 = loader.CMAPSSLoader(3)
172 | 
173 |     def test_get_pair_idx_piecewise(self):
174 |         data = datasets.cmapss.PairedCMAPSS([self.cmapss_normal], "dev", 512, 1, True)
175 |         middle_idx = self.length // 2
176 |         for _ in range(512):
177 |             run, anchor_idx, query_idx, distance, _ = data._get_pair_idx_piecewise()
178 |             self.assertEqual(middle_idx, len(run) // 2)
179 |             if anchor_idx < middle_idx:
180 |                 self.assertEqual(0, distance)
181 |             else:
182 |                 self.assertLessEqual(0, distance)
183 | 
184 |     def test_get_pair_idx_linear(self):
185 |         data = datasets.cmapss.PairedCMAPSS([self.cmapss_normal], "dev", 512, 1, True)
186 |         for _ in range(512):
187 |             run, anchor_idx, query_idx, distance, _ = data._get_pair_idx()
188 |             self.assertLess(0, distance)
189 |             self.assertGreaterEqual(125, distance)
190 | 
191 |     def test_get_labeled_pair_idx(self):
192 |         data = datasets.cmapss.PairedCMAPSS([self.cmapss_normal], "dev", 512, 1, True)
193 |         for _ in range(512):
194 |             run, anchor_idx, query_idx, distance, _ = data._get_labeled_pair_idx()
195 |             self.assertEqual(self.length, len(run))
196 |             expected_distance = data._labels[0][anchor_idx] - data._labels[0][query_idx]
197 |             self.assertLessEqual(0, distance)
198 |             self.assertEqual(expected_distance, distance)
199 | 
200 |     def test_pair_func_selection(self):
201 |         with self.subTest("default"):
202 |             data = datasets.cmapss.PairedCMAPSS([self.cmapss_normal], "dev", 512, 1, True)
203 |             self.assertEqual(data._get_pair_idx, data._get_pair_func)
204 |         with self.subTest("piecewise"):
205 |             data = datasets.cmapss.PairedCMAPSS(
206 |                 [self.cmapss_normal], "dev", 512, 1, True, mode="linear"
207 |             )
208 |             self.assertEqual(data._get_pair_idx, data._get_pair_func)
209 |         with self.subTest("piecewise"):
210 |             data = datasets.cmapss.PairedCMAPSS(
211 |                 [self.cmapss_normal], "dev", 512, 1, True, mode="piecewise"
212 |             )
213 |             self.assertEqual(data._get_pair_idx_piecewise, data._get_pair_func)
214 |         with self.subTest("labeled"):
215 |             data = datasets.cmapss.PairedCMAPSS(
216 |                 [self.cmapss_normal], "dev", 512, 1, True, mode="labeled"
217 |             )
218 |             self.assertEqual(data._get_labeled_pair_idx, data._get_pair_func)
219 | 
220 |     def test_sampled_data(self):
221 |         fixed_idx = [0, 60, 80, 1, 55, 99]  # two samples with run, anchor and query idx
222 |         data = datasets.cmapss.PairedCMAPSS([self.cmapss_short], "dev", 2, 2)
223 |         data._rng = mock.MagicMock()
224 |         data._rng.integers = mock.MagicMock(side_effect=fixed_idx)
225 |         for i, sample in enumerate(data):
226 |             idx = 3 * i
227 |             expected_run = data._features[fixed_idx[idx]]
228 |             expected_anchor = expected_run[fixed_idx[idx + 1]]
229 |             expected_query = expected_run[fixed_idx[idx + 2]]
230 |             expected_distance = min(125, fixed_idx[idx + 2] - fixed_idx[idx + 1]) / 125
231 |             expected_domain_idx = 0
232 |             self.assertEqual(0, torch.dist(expected_anchor, sample[0]))
233 |             self.assertEqual(0, torch.dist(expected_query, sample[1]))
234 |             self.assertAlmostEqual(expected_distance, sample[2].item())
235 |             self.assertEqual(expected_domain_idx, sample[3].item())
236 | 
237 |     def test_discarding_too_short_runs(self):
238 |         data = datasets.cmapss.PairedCMAPSS([self.cmapss_short], "dev", 512, 2)
239 |         for run, labels in zip(data._features, data._labels):
240 |             self.assertTrue((run >= 1).all())
241 |             self.assertTrue((labels < 500).all())
242 |             self.assertLess(2, len(run))
243 | 
244 |     def test_determinisim(self):
245 |         with self.subTest("non-deterministic"):
246 |             data = datasets.cmapss.PairedCMAPSS([self.cmapss_short], "dev", 512, 2)
247 |             self.assertTrue(self._two_epochs_different(data))
248 |         with self.subTest("non-deterministic"):
249 |             data = datasets.cmapss.PairedCMAPSS(
250 |                 [self.cmapss_short], "dev", 512, 2, deterministic=True
251 |             )
252 |             self.assertFalse(self._two_epochs_different(data))
253 | 
254 |     def _two_epochs_different(self, data):
255 |         first_epoch = [samples for samples in data]
256 |         second_epoch = [samples for samples in data]
257 |         different = False
258 |         for first_samples, second_samples in zip(first_epoch, second_epoch):
259 |             for first, second in zip(first_samples, second_samples):
260 |                 if torch.dist(first, second) > 0:
261 |                     different = True
262 | 
263 |         return different
264 | 
265 |     def test_min_distance(self):
266 |         dataset = datasets.cmapss.PairedCMAPSS(
267 |             [self.cmapss_short], "dev", 512, min_distance=30
268 |         )
269 |         pairs = self._get_pairs(dataset)
270 |         distances = pairs[:, 1] - pairs[:, 0]
271 |         self.assertTrue(np.all(pairs[:, 0] < pairs[:, 1]))
272 |         self.assertTrue(np.all(distances >= 30))
273 | 
274 |     def test_build_pairs(self):
275 |         for split in ["dev", "val"]:
276 |             with self.subTest(split=split):
277 |                 paired_dataset = datasets.cmapss.PairedCMAPSS(
278 |                     [self.fd1, self.fd3], split, 1000, 1
279 |                 )
280 |                 pairs = self._get_pairs(paired_dataset)
281 |                 self.assertTrue(
282 |                     np.all(pairs[:, 0] < pairs[:, 1])
283 |                 )  # query always after anchor
284 |                 self.assertTrue(np.all(pairs[:, 3] <= 1))  # domain label is either 1
285 |                 self.assertTrue(np.all(pairs[:, 3] >= 0))  # or zero
286 | 
287 |     def test_domain_labels(self):
288 |         dataset = datasets.cmapss.PairedCMAPSS(
289 |             [self.cmapss_normal, self.cmapss_short], "dev", 512, min_distance=30
290 |         )
291 |         for _ in range(512):
292 |             run, _, _, _, domain_idx = dataset._get_pair_idx_piecewise()
293 |             if len(run) == self.length:
294 |                 self.assertEqual(0, domain_idx)  # First domain is self.length long
295 |             else:
296 |                 self.assertEqual(1, domain_idx)  # Second is not
297 | 
298 |     def _get_pairs(self, paired):
299 |         pairs = [paired._get_pair_idx_piecewise() for _ in range(paired.num_samples)]
300 |         pairs = [
301 |             (anchor_idx, query_idx, distance, domain_idx)
302 |             for _, anchor_idx, query_idx, distance, domain_idx in pairs
303 |         ]
304 | 
305 |         return np.array(pairs)
306 | 
307 | 
308 | class TestAdaptionDataset(unittest.TestCase):
309 |     def setUp(self):
310 |         self.source = TensorDataset(torch.arange(100), torch.arange(100))
311 |         self.target = TensorDataset(torch.arange(150), torch.arange(150))
312 |         self.dataset = datasets.cmapss.AdaptionDataset(
313 |             self.source,
314 |             self.target,
315 |         )
316 | 
317 |     def test_len(self):
318 |         self.assertEqual(len(self.dataset.source), len(self.dataset))
319 | 
320 |     def test_source_target_shuffeled(self):
321 |         for i in range(len(self.dataset)):
322 |             source_one, label_one, target_one = self.dataset[i]
323 |             source_another, label_another, target_another = self.dataset[i]
324 |             self.assertEqual(source_one, source_another)
325 |             self.assertEqual(label_one, label_another)
326 |             self.assertNotEqual(target_one, target_another)
327 | 
328 |     def test_source_target_deterministic(self):
329 |         dataset = datasets.cmapss.AdaptionDataset(
330 |             self.source, self.target, deterministic=True
331 |         )
332 |         for i in range(len(dataset)):
333 |             source_one, label_one, target_one = dataset[i]
334 |             source_another, label_another, target_another = dataset[i]
335 |             self.assertEqual(source_one, source_another)
336 |             self.assertEqual(label_one, label_another)
337 |             self.assertEqual(target_one, target_another)
338 | 
339 |     def test_source_sampled_completely(self):
340 |         for i in range(len(self.dataset)):
341 |             source, labels, _ = self.dataset[i]
342 |             self.assertEqual(i, source.item())
343 |             self.assertEqual(i, labels.item())
344 | 


--------------------------------------------------------------------------------
/tests/dataset_tests/test_loader.py:
--------------------------------------------------------------------------------
  1 | import unittest
  2 | from unittest import mock
  3 | 
  4 | import torch
  5 | 
  6 | from datasets import loader
  7 | 
  8 | 
  9 | class TestCMAPSSLoader(unittest.TestCase):
 10 |     NUM_CHANNELS = len(loader.CMAPSSLoader.DEFAULT_CHANNELS)
 11 | 
 12 |     def test_run_shape_and_dtype(self):
 13 |         window_sizes = [30, 20, 30, 15]
 14 |         for n, win in enumerate(window_sizes, start=1):
 15 |             cmapss_loader = loader.CMAPSSLoader(n)
 16 |             for split in ["dev", "val", "test"]:
 17 |                 with self.subTest(fd=n, split=split):
 18 |                     self._check_split(cmapss_loader, split, win)
 19 | 
 20 |     def _check_split(self, cmapss_loader, split, window_size):
 21 |         features, targets = cmapss_loader.load_split(split)
 22 |         for run, run_target in zip(features, targets):
 23 |             self._assert_run_correct(run, run_target, window_size)
 24 | 
 25 |     def _assert_run_correct(self, run, run_target, win):
 26 |         self.assertEqual(win, run.shape[2])
 27 |         self.assertEqual(self.NUM_CHANNELS, run.shape[1])
 28 |         self.assertEqual(len(run), len(run_target))
 29 |         self.assertEqual(torch.float32, run.dtype)
 30 |         self.assertEqual(torch.float32, run_target.dtype)
 31 | 
 32 |     def test_override_window_size(self):
 33 |         window_size = 40
 34 |         for n in range(1, 5):
 35 |             dataset = loader.CMAPSSLoader(n, window_size=window_size)
 36 |             for split in ["dev", "val", "test"]:
 37 |                 with self.subTest(fd=n, split=split):
 38 |                     features, targets = dataset.load_split(split)
 39 |                     for run in features:
 40 |                         self.assertEqual(window_size, run.shape[2])
 41 | 
 42 |     def test_feature_select(self):
 43 |         dataset = loader.CMAPSSLoader(1, feature_select=[4, 9, 10, 13, 14, 15, 22])
 44 |         for split in ["dev", "val", "test"]:
 45 |             features, _ = dataset.load_split(split)
 46 |             for run in features:
 47 |                 self.assertEqual(7, run.shape[1])
 48 | 
 49 |     def test_percent_fail_runs(self):
 50 |         full_dataset = loader.CMAPSSLoader(fd=1, window_size=30)
 51 |         full_dev, full_dev_targets = full_dataset.load_split("dev")
 52 |         full_val = full_dataset.load_split("val")[0]
 53 |         full_test = full_dataset.load_split("test")[0]
 54 | 
 55 |         dataset = loader.CMAPSSLoader(fd=1, window_size=30, percent_fail_runs=0.8)
 56 |         trunc_dev, trunc_dev_targets = dataset.load_split("dev")
 57 | 
 58 |         self.assertGreater(len(full_dev), len(trunc_dev))
 59 |         self.assertAlmostEqual(0.8, len(trunc_dev) / len(full_dev), delta=0.01)
 60 |         self.assertEqual(len(full_val), len(dataset.load_split("val")[0]))
 61 |         self.assertEqual(len(full_test), len(dataset.load_split("test")[0]))
 62 | 
 63 |         for full_run, trunc_run in zip(full_dev, trunc_dev):
 64 |             self.assertEqual(0.0, torch.dist(trunc_run, full_run))
 65 |         for full_targets, trunc_targets in zip(full_dev_targets, trunc_dev_targets):
 66 |             self.assertEqual(0.0, torch.dist(trunc_targets, full_targets))
 67 | 
 68 |     def test_percent_broken(self):
 69 |         full_dataset = loader.CMAPSSLoader(fd=1, window_size=30)
 70 |         full_dev, full_dev_targets = full_dataset.load_split("dev")
 71 |         full_val = full_dataset.load_split("val")[0]
 72 |         full_test = full_dataset.load_split("test")[0]
 73 | 
 74 |         dataset = loader.CMAPSSLoader(fd=1, window_size=30, percent_broken=0.2)
 75 |         truncated_dev, truncated_dev_targets = dataset.load_split("dev")
 76 | 
 77 |         full_length = sum(len(r) for r in full_dev)
 78 |         truncated_length = sum(len(r) for r in truncated_dev)
 79 |         self.assertGreater(full_length, truncated_length)
 80 |         self.assertAlmostEqual(0.2, truncated_length / full_length, delta=0.01)
 81 | 
 82 |         # No failure data in truncated dev data
 83 |         self.assertFalse(any(torch.any(r == 1) for r in truncated_dev_targets))
 84 | 
 85 |         for full_run, trunc_run in zip(full_dev, truncated_dev):
 86 |             trunc_run_length = trunc_run.shape[0]
 87 |             self.assertEqual(0.0, torch.dist(trunc_run, full_run[:trunc_run_length]))
 88 |         for full_run, trunc_run in zip(full_dev_targets, truncated_dev_targets):
 89 |             trunc_run_length = trunc_run.shape[0]
 90 |             self.assertEqual(0.0, torch.dist(trunc_run, full_run[:trunc_run_length]))
 91 | 
 92 |         # Val dataset_tests not truncated
 93 |         for full_run, trunc_run in zip(full_val, dataset.load_split("val")[0]):
 94 |             self.assertEqual(len(full_run), len(trunc_run))
 95 | 
 96 |         # Test dataset_tests not truncated
 97 |         for full_run, trunc_run in zip(full_test, dataset.load_split("test")[0]):
 98 |             self.assertEqual(len(full_run), len(trunc_run))
 99 | 
100 |     @mock.patch(
101 |         "datasets.loader.CMAPSSLoader._truncate_runs", wraps=lambda x, y, *args: (x, y)
102 |     )
103 |     def test_val_truncation(self, mock_truncate):
104 |         dataset = loader.CMAPSSLoader(fd=1, window_size=30)
105 |         with self.subTest(truncate_val=False):
106 |             dataset.load_split("dev")
107 |             mock_truncate.assert_called_once()
108 |             mock_truncate.reset_mock()
109 |             dataset.load_split("val")
110 |             mock_truncate.assert_not_called()
111 | 
112 |         dataset = loader.CMAPSSLoader(fd=1, window_size=30, truncate_val=True)
113 |         with self.subTest(truncate_val=True):
114 |             dataset.load_split("dev")
115 |             mock_truncate.assert_called_once()
116 |             mock_truncate.reset_mock()
117 |             dataset.load_split("val")
118 |             mock_truncate.assert_called_once()
119 | 
120 |     def test_normalization_min_max(self):
121 |         for i in range(1, 5):
122 |             with self.subTest(fd=i):
123 |                 full_dataset = loader.CMAPSSLoader(fd=i, window_size=30)
124 |                 full_dev, full_dev_targets = full_dataset.load_split("dev")
125 | 
126 |                 self.assertAlmostEqual(max(torch.max(r).item() for r in full_dev), 1.0)
127 |                 self.assertAlmostEqual(min(torch.min(r).item() for r in full_dev), -1.0)
128 | 
129 |                 truncated_dataset = loader.CMAPSSLoader(
130 |                     fd=i, window_size=30, percent_fail_runs=0.8
131 |                 )
132 |                 trunc_dev, trunc_dev_targets = truncated_dataset.load_split("dev")
133 |                 self.assertLessEqual(max(torch.max(r).item() for r in trunc_dev), 1.0)
134 |                 self.assertGreaterEqual(min(torch.min(r).item() for r in trunc_dev), -1.0)
135 | 
136 |                 truncated_dataset = loader.CMAPSSLoader(
137 |                     fd=i, window_size=30, percent_broken=0.2
138 |                 )
139 |                 trunc_dev, trunc_dev_targets = truncated_dataset.load_split("dev")
140 |                 self.assertLessEqual(max(torch.max(r).item() for r in trunc_dev), 1.0)
141 |                 self.assertGreaterEqual(min(torch.min(r).item() for r in trunc_dev), -1.0)
142 | 
143 |     def test_truncation_by_index(self):
144 |         full_dataset = loader.CMAPSSLoader(1)
145 |         full_train, full_train_targets = full_dataset.load_split("dev")
146 |         indices = [50, 60, 70]
147 |         truncated_dataset = loader.CMAPSSLoader(1, percent_fail_runs=indices)
148 |         trunc_train, trunc_train_targets = truncated_dataset.load_split("dev")
149 | 
150 |         self.assertEqual(len(indices), len(trunc_train))
151 |         self.assertEqual(len(indices), len(trunc_train_targets))
152 |         for trunc_idx, full_idx in enumerate(indices):
153 |             self.assertEqual(0, torch.dist(full_train[full_idx], trunc_train[trunc_idx]))
154 |             self.assertEqual(
155 |                 0,
156 |                 torch.dist(full_train_targets[full_idx], trunc_train_targets[trunc_idx]),
157 |             )
158 | 


--------------------------------------------------------------------------------
/tests/lightning_tests/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/tilman151/self-supervised-ssl/745d4b8b06a946f1be1e2c8bc15e4b0b55c7c4ff/tests/lightning_tests/__init__.py


--------------------------------------------------------------------------------
/tests/lightning_tests/test_autoencoder.py:
--------------------------------------------------------------------------------
  1 | import unittest
  2 | from unittest import mock
  3 | 
  4 | import torch
  5 | 
  6 | from lightning import autoencoder
  7 | 
  8 | 
  9 | class TestUnsupervisedPretraining(unittest.TestCase):
 10 |     def setUp(self):
 11 |         self.net = autoencoder.AutoencoderPretraining(
 12 |             in_channels=14,
 13 |             seq_len=30,
 14 |             num_layers=4,
 15 |             kernel_size=3,
 16 |             base_filters=16,
 17 |             latent_dim=64,
 18 |             dropout=0.1,
 19 |             domain_tradeoff=0.0,
 20 |             domain_disc_dim=32,
 21 |             num_disc_layers=2,
 22 |             weight_decay=0,
 23 |             lr=0.01,
 24 |         )
 25 | 
 26 |     def test_mode_hparam(self):
 27 |         self.assertLess(1, len(self.net.hparams))
 28 |         self.assertIn("mode", self.net.hparams)
 29 |         self.assertEqual("autoencoder", self.net.hparams["mode"])
 30 | 
 31 |     @torch.no_grad()
 32 |     def test_encoder(self):
 33 |         inputs = torch.randn(16, 14, 30)
 34 |         outputs = self.net.encoder(inputs)
 35 |         self.assertEqual(torch.Size((16, 64)), outputs.shape)
 36 | 
 37 |     @torch.no_grad()
 38 |     def test_decoder(self):
 39 |         inputs = torch.randn(16, 64)
 40 |         outputs = self.net.decoder(inputs)
 41 |         self.assertEqual(torch.Size((16, 14, 30)), outputs.shape)
 42 | 
 43 |     @torch.no_grad()
 44 |     def test_forward(self):
 45 |         inputs = torch.randn(16, 14, 30)
 46 |         outputs = self.net(inputs)
 47 |         self.assertEqual(torch.Size((16, 14, 30)), outputs.shape)
 48 | 
 49 |     def test_batch_independence(self):
 50 |         torch.autograd.set_detect_anomaly(True)
 51 | 
 52 |         inputs = torch.randn(16, 14, 30)
 53 |         inputs.requires_grad = True
 54 | 
 55 |         # Compute forward pass in eval mode to deactivate batch norm
 56 |         self.net.eval()
 57 |         outputs = self.net(inputs)
 58 |         self.net.train()
 59 | 
 60 |         # Mask loss for certain samples in batch
 61 |         batch_size = outputs.shape[0]
 62 |         mask_idx = torch.randint(0, batch_size, ())
 63 |         mask = torch.ones_like(outputs)
 64 |         mask[mask_idx] = 0
 65 |         output = outputs * mask
 66 | 
 67 |         # Compute backward pass
 68 |         loss = output.mean()
 69 |         loss.backward(retain_graph=True)
 70 | 
 71 |         # Check if gradient exists and is zero for masked samples
 72 |         for i, grad in enumerate(inputs.grad[:batch_size]):
 73 |             if i == mask_idx:
 74 |                 self.assertTrue(torch.all(grad == 0).item())
 75 |             else:
 76 |                 self.assertTrue(not torch.all(grad == 0))
 77 |         inputs.grad = None
 78 | 
 79 |         torch.autograd.set_detect_anomaly(False)
 80 | 
 81 |     def test_all_parameters_updated(self):
 82 |         optim = torch.optim.SGD(self.net.parameters(), lr=0.1)
 83 | 
 84 |         inputs = (
 85 |             torch.randn(16, 14, 30),
 86 |             torch.randn(16, 14, 30),
 87 |             torch.randn(16),
 88 |             torch.randn(16),
 89 |         )
 90 |         loss = self.net.training_step(inputs, batch_idx=0)
 91 |         loss.backward()
 92 |         optim.step()
 93 | 
 94 |         for param_name, param in self.net.named_parameters():
 95 |             if param.requires_grad:
 96 |                 with self.subTest(name=param_name):
 97 |                     self.assertIsNotNone(param.grad)
 98 |                     self.assertNotEqual(0.0, torch.sum(param.grad ** 2))
 99 | 
100 |     @mock.patch("pytorch_lightning.LightningModule.log")
101 |     @torch.no_grad()
102 |     def test_metric_val_reduction(self, mock_log):
103 |         num_batches = 600
104 |         embeddings = torch.randn(num_batches, 64, 64)
105 |         reconstructions = torch.randn(num_batches, 64, 14, 30)
106 |         domain_predictions = torch.randn(num_batches, 32)
107 |         self._mock_predictions(embeddings, reconstructions, domain_predictions)
108 | 
109 |         self._feed_dummy_val(num_batches)
110 | 
111 |         regression_loss = torch.mean(reconstructions ** 2)
112 |         domain_loss = torch.tensor(0.0)
113 |         checkpoint_score = regression_loss - 0.1 * domain_loss
114 |         expected_logs = {
115 |             "val/regression_loss": (regression_loss, 0.0001),
116 |             "val/domain_loss": (domain_loss, 0.0001),
117 |             "val/checkpoint_score": (checkpoint_score, 0.0001),
118 |         }
119 |         self.net.validation_epoch_end([])
120 | 
121 |         self._assert_logs(mock_log, expected_logs)
122 | 
123 |     def test_metric_val_updates(self):
124 |         num_batches = 600
125 |         embeddings = torch.randn(num_batches, 64, 64)
126 |         reconstructions = torch.randn(num_batches, 64, 14, 30)
127 |         domain_predictions = torch.randn(num_batches, 32)
128 |         self._mock_predictions(embeddings, reconstructions, domain_predictions)
129 | 
130 |         self._feed_dummy_val(num_batches)
131 |         self.assertEqual(num_batches, self.net.regression_metric.sample_counter)
132 |         self.assertEqual(num_batches, self.net.domain_metric.sample_counter)
133 | 
134 |     def _mock_predictions(self, embeddings, reconstructions, domain_predictions):
135 |         self.net.encoder.forward = mock.MagicMock(side_effect=embeddings)
136 |         self.net.decoder.forward = mock.MagicMock(side_effect=reconstructions)
137 | 
138 |     def _feed_dummy_val(self, num_batches):
139 |         for i in range(num_batches):
140 |             self.net.validation_step(
141 |                 (
142 |                     torch.zeros(32, 14, 30),
143 |                     torch.zeros(32, 14, 30),
144 |                     torch.zeros(32),
145 |                     torch.zeros(32),
146 |                 ),
147 |                 batch_idx=i,
148 |                 dataloader_idx=0,
149 |             )
150 | 
151 |     def _assert_logs(self, mock_log, expected_logs):
152 |         mock_log.assert_called()
153 |         for call in mock_log.mock_calls:
154 |             metric = call[1][0]
155 |             self.assertIn(metric, expected_logs, "Unexpected logged metric found.")
156 |             expected_value, delta = expected_logs[metric]
157 |             expected_value = expected_value.item()
158 |             actual_value = call[1][1].item()
159 |             with self.subTest(metric):
160 |                 self.assertAlmostEqual(expected_value, actual_value, delta=delta)
161 | 


--------------------------------------------------------------------------------
/tests/lightning_tests/test_baseline.py:
--------------------------------------------------------------------------------
  1 | import unittest
  2 | from unittest import mock
  3 | 
  4 | import torch
  5 | 
  6 | import building
  7 | from lightning import baseline
  8 | 
  9 | 
 10 | class TestBaselineTemplate:
 11 |     @torch.no_grad()
 12 |     def test_encoder(self):
 13 |         inputs = torch.randn(16, 14, 30)
 14 |         outputs = self.net.encoder(inputs)
 15 |         self.assertEqual(self.encoder_shape, outputs.shape)
 16 | 
 17 |     @torch.no_grad()
 18 |     def test_regressor(self):
 19 |         inputs = torch.randn(*self.encoder_shape)
 20 |         outputs = self.net.regressor(inputs)
 21 |         self.assertEqual(torch.Size((16,)), outputs.shape)
 22 | 
 23 |     def test_batch_independence(self):
 24 |         inputs = torch.randn(16, 14, 30)
 25 |         inputs.requires_grad = True
 26 | 
 27 |         # Compute forward pass in eval mode to deactivate batch norm
 28 |         self.net.eval()
 29 |         outputs = self.net(inputs)
 30 |         self.net.train()
 31 | 
 32 |         # Mask loss for certain samples in batch
 33 |         batch_size = outputs.shape[0]
 34 |         mask_idx = torch.randint(0, batch_size, ())
 35 |         mask = torch.ones_like(outputs)
 36 |         mask[mask_idx] = 0
 37 |         output = outputs * mask
 38 | 
 39 |         # Compute backward pass
 40 |         loss = output.mean()
 41 |         loss.backward(retain_graph=True)
 42 | 
 43 |         # Check if gradient exists and is zero for masked samples
 44 |         for i, grad in enumerate(inputs.grad[:batch_size]):
 45 |             if i == mask_idx:
 46 |                 self.assertTrue(torch.all(grad == 0).item())
 47 |             else:
 48 |                 self.assertTrue(not torch.all(grad == 0))
 49 |         inputs.grad = None
 50 | 
 51 |     def test_all_parameters_updated(self):
 52 |         optim = torch.optim.SGD(self.net.parameters(), lr=0.1)
 53 | 
 54 |         loss = self.net.training_step(
 55 |             (torch.randn(16, 14, 30), torch.ones(16)), batch_idx=0
 56 |         )
 57 |         loss.backward()
 58 |         optim.step()
 59 | 
 60 |         for param_name, param in self.net.named_parameters():
 61 |             if param.requires_grad:
 62 |                 with self.subTest(name=param_name):
 63 |                     self.assertIsNotNone(param.grad)
 64 |                     self.assertNotEqual(0.0, torch.sum(param.grad ** 2))
 65 | 
 66 |     @mock.patch("pytorch_lightning.LightningModule.log")
 67 |     @torch.no_grad()
 68 |     def test_metric_val_reduction(self, mock_log):
 69 |         source_batches = 600
 70 |         source_prediction = torch.randn(source_batches, 32) + 30
 71 |         self._mock_predictions(source_prediction)
 72 | 
 73 |         self._feed_dummy_val(source_batches)
 74 | 
 75 |         source_regression_loss = torch.sqrt(torch.mean(source_prediction ** 2))
 76 |         expected_logs = {
 77 |             "val/regression_loss": (source_regression_loss, 0.001),
 78 |         }
 79 |         self.net.validation_epoch_end([])
 80 | 
 81 |         self._assert_logs(mock_log, expected_logs)
 82 | 
 83 |     def test_metric_val_updates(self):
 84 |         source_batches = 600
 85 |         source_prediction = torch.randn(source_batches, 32) + 30
 86 |         self._mock_predictions(source_prediction)
 87 | 
 88 |         self._feed_dummy_val(source_batches)
 89 |         self.assertEqual(source_batches, self.net.regression_metrics[1].sample_counter)
 90 | 
 91 |     @mock.patch("pytorch_lightning.LightningModule.log")
 92 |     @torch.no_grad()
 93 |     def test_metric_test_reduction(self, mock_log):
 94 |         source_batches = 600
 95 |         source_prediction = [
 96 |             torch.randn(source_batches, 32) + i * 10 for i in range(1, 5)
 97 |         ]
 98 |         self._mock_predictions(torch.cat(source_prediction))
 99 | 
100 |         self._feed_dummy_test(source_batches)
101 | 
102 |         expected_logs = {}
103 |         for i in range(4):
104 |             source_regression_loss = torch.sqrt(torch.mean(source_prediction[i] ** 2))
105 |             metric_name = f"test/regression_loss_fd{i+1}"
106 |             expected_logs[metric_name] = (source_regression_loss, 0.001)
107 |         self.net.test_epoch_end([])
108 | 
109 |         self._assert_logs(mock_log, expected_logs)
110 | 
111 |     def test_metric_test_updates(self):
112 |         source_batches = 600
113 |         source_prediction = [
114 |             torch.randn(source_batches, 32) + i * 10 for i in range(1, 5)
115 |         ]
116 |         self._mock_predictions(torch.cat(source_prediction))
117 | 
118 |         self._feed_dummy_test(source_batches)
119 |         for fd in range(1, 5):
120 |             self.assertEqual(
121 |                 source_batches, self.net.regression_metrics[fd].sample_counter
122 |             )
123 | 
124 |     def _mock_predictions(self, prediction):
125 |         self.net.regressor.forward = mock.MagicMock(side_effect=prediction)
126 | 
127 |     def _feed_dummy_val(self, num_batches):
128 |         for i in range(num_batches):
129 |             self.net.validation_step(
130 |                 (torch.zeros(32, 14, 30), torch.zeros(32)),
131 |                 batch_idx=i,
132 |             )
133 | 
134 |     def _feed_dummy_test(self, num_batches):
135 |         for fd in range(4):
136 |             for i in range(num_batches):
137 |                 self.net.test_step(
138 |                     (torch.zeros(32, 14, 30), torch.zeros(32)),
139 |                     batch_idx=i,
140 |                     dataloader_idx=fd,
141 |                 )
142 | 
143 |     def _assert_logs(self, mock_log, expected_logs):
144 |         mock_log.assert_called()
145 |         for call in mock_log.mock_calls:
146 |             metric = call[1][0]
147 |             self.assertIn(metric, expected_logs, "Unexpected logged metric found.")
148 |             expected_value, delta = expected_logs[metric]
149 |             expected_value = expected_value.item()
150 |             actual_value = call[1][1].item()
151 |             with self.subTest(metric):
152 |                 self.assertAlmostEqual(expected_value, actual_value, delta=delta)
153 | 
154 | 
155 | class TestCnnBaseline(unittest.TestCase, TestBaselineTemplate):
156 |     def setUp(self):
157 |         self.encoder_shape = torch.Size((16, 64))
158 |         self.trade_off = 0.5
159 |         self.net = baseline.Baseline(
160 |             in_channels=14,
161 |             seq_len=30,
162 |             num_layers=4,
163 |             kernel_size=3,
164 |             base_filters=16,
165 |             latent_dim=64,
166 |             dropout=0.1,
167 |             optim_type="adam",
168 |             lr=0.01,
169 |         )
170 | 
171 |     def test_load_from_rbm(self):
172 |         self.net.encoder.layers[1].bias.data = torch.randn_like(
173 |             self.net.encoder.layers[1].bias.data
174 |         )
175 |         old_weight_shape = self.net.encoder.layers[0].weight.shape
176 |         old_bias_shape = self.net.encoder.layers[1].bias.shape
177 |         rbm = building.build_rbm(14, 16)
178 |         self.assertNotEqual(
179 |             0,
180 |             torch.dist(rbm.interaction.module.weight, self.net.encoder.layers[0].weight),
181 |         )
182 |         self.assertNotEqual(
183 |             0, torch.dist(rbm.hidden.bias, self.net.encoder.layers[1].bias)
184 |         )
185 |         self.assertTrue(self.net.encoder.layers[0].weight.requires_grad)
186 |         self.assertTrue(self.net.encoder.layers[1].bias.requires_grad)
187 | 
188 |         self.net.load_from_rbm(rbm)
189 |         self.assertEqual(
190 |             0,
191 |             torch.dist(rbm.interaction.module.weight, self.net.encoder.layers[0].weight),
192 |         )
193 |         self.assertEqual(0, torch.dist(rbm.hidden.bias, self.net.encoder.layers[1].bias))
194 |         self.assertTrue(self.net.encoder.layers[0].weight.requires_grad)
195 |         self.assertTrue(self.net.encoder.layers[1].bias.requires_grad)
196 |         self.assertEqual(old_weight_shape, self.net.encoder.layers[0].weight.shape)
197 |         self.assertEqual(old_bias_shape, self.net.encoder.layers[1].bias.shape)
198 | 


--------------------------------------------------------------------------------
/tests/lightning_tests/test_loggers.py:
--------------------------------------------------------------------------------
 1 | import os
 2 | import shutil
 3 | import tempfile
 4 | import unittest
 5 | from unittest import mock
 6 | 
 7 | import pytorch_lightning as pl
 8 | 
 9 | import datasets
10 | import lightning.loggers as loggers
11 | from lightning import baseline
12 | 
13 | 
14 | class TestMLTBLogger(unittest.TestCase):
15 |     def setUp(self):
16 |         self.logdir = tempfile.mkdtemp()
17 |         self.test_tag = "test_tag"
18 |         self.logger = loggers.MLTBLogger(
19 |             self.logdir, "Test", self.test_tag, {"pb": 0.1, "foo": "bar"}
20 |         )
21 | 
22 |     def test_save_path(self):
23 |         self._run_dummy_training()
24 |         artifact_path = self.logger.checkpoint_path
25 |         self.assertListEqual(["epoch=0.ckpt"], os.listdir(artifact_path))
26 | 
27 |     def test_tensorboard_struct(self):
28 |         self._run_dummy_training()
29 |         expected_tf_logdir = os.path.join(
30 |             self.logdir, "tensorboard", "Test", "0.1pb", "foo-bar", "version_0"
31 |         )
32 |         self.assertTrue(os.path.exists(expected_tf_logdir))
33 | 
34 |     def test_version_tag(self):
35 |         self._run_dummy_training()
36 |         run_id = self.logger._mlflow_logger.run_id
37 |         run = self.logger.mlflow_experiment.get_run(run_id)
38 |         self.assertDictEqual({"version": "test_tag"}, run.data.tags)
39 | 
40 |     def _run_dummy_training(self):
41 |         trainer = pl.Trainer(
42 |             gpus=0,
43 |             max_epochs=1,
44 |             logger=self.logger,
45 |             deterministic=True,
46 |             log_every_n_steps=10,
47 |         )
48 |         data = datasets.BaselineDataModule(fd_source=1, batch_size=512, window_size=30)
49 |         model = baseline.Baseline(
50 |             in_channels=14,
51 |             seq_len=30,
52 |             num_layers=1,
53 |             kernel_size=1,
54 |             base_filters=1,
55 |             latent_dim=16,
56 |             optim_type="adam",
57 |             lr=0.01,
58 |             record_embeddings=False,
59 |         )
60 |         trainer.fit(model, datamodule=data)
61 | 
62 |     def tearDown(self):
63 |         shutil.rmtree(self.logdir)
64 | 
65 | 
66 | class TestMinEpochModelCheckpoint(unittest.TestCase):
67 |     @mock.patch("pytorch_lightning.callbacks.ModelCheckpoint.save_checkpoint")
68 |     def test_min_epoch(self, mock_super_save_checkpoint):
69 |         mock_trainer = mock.MagicMock("trainer")
70 |         mock_pl_module = mock.MagicMock("pl_module")
71 | 
72 |         min_epoch = 1
73 |         checkpoint_callback = loggers.MinEpochModelCheckpoint(
74 |             "val/test_metric", min_epochs_before_saving=min_epoch
75 |         )
76 | 
77 |         with self.subTest(case="before_min_epoch"):
78 |             mock_trainer.current_epoch = 0
79 |             checkpoint_callback.save_checkpoint(mock_trainer, mock_pl_module)
80 |             mock_super_save_checkpoint.assert_not_called()
81 | 
82 |         with self.subTest(case="after_min_epoch"):
83 |             mock_trainer.current_epoch = 2
84 |             checkpoint_callback.save_checkpoint(mock_trainer, mock_pl_module)
85 |             mock_super_save_checkpoint.assert_called_with(mock_trainer, mock_pl_module)
86 | 


--------------------------------------------------------------------------------
/tests/lightning_tests/test_metrics.py:
--------------------------------------------------------------------------------
  1 | import unittest
  2 | 
  3 | import matplotlib.pyplot as plt
  4 | import torch
  5 | 
  6 | from lightning import metrics
  7 | 
  8 | 
  9 | class TestEmbeddingViz(unittest.TestCase):
 10 |     def setUp(self):
 11 |         self.metric = metrics.EmbeddingViz(32, 128)
 12 | 
 13 |     def test_updating(self):
 14 |         embeddings = torch.randn(16, 128)
 15 |         labels = torch.ones(16)
 16 |         ruls = torch.arange(0, 16)
 17 | 
 18 |         self.metric.update(embeddings, labels, ruls)
 19 |         self.metric.update(embeddings, labels, ruls)
 20 | 
 21 |         self.assertEqual(0, torch.sum(self.metric.embeddings[:16] - embeddings))
 22 |         self.assertEqual(0, torch.sum(self.metric.labels[:16] - labels))
 23 |         self.assertEqual(0, torch.sum(self.metric.embeddings[16:32] - embeddings))
 24 |         self.assertEqual(0, torch.sum(self.metric.labels[16:32] - labels))
 25 | 
 26 |     def test_compute(self):
 27 |         embeddings = torch.randn(32, 128)
 28 |         labels = torch.cat([torch.ones(16), torch.zeros(16)])
 29 |         ruls = torch.arange(0, 32)
 30 | 
 31 |         self.metric.update(embeddings, labels, ruls)
 32 |         fig = self.metric.compute()
 33 | 
 34 |         self.assertIsInstance(fig, plt.Figure)
 35 |         self.assertListEqual([20.0, 10.0], list(fig.get_size_inches()))
 36 | 
 37 | 
 38 | class TestRULScore(unittest.TestCase):
 39 |     def setUp(self):
 40 |         self.metric = metrics.RULScore()
 41 | 
 42 |     def test_negative(self):
 43 |         inputs = torch.ones(2)
 44 |         targets = torch.ones(2) * 2
 45 |         actual_score = self.metric(inputs, targets)
 46 |         expected_score = torch.sum(
 47 |             torch.exp((inputs - targets) / self.metric.neg_factor) - 1
 48 |         )
 49 |         self.assertEqual(expected_score, actual_score)
 50 | 
 51 |     def test_positive(self):
 52 |         inputs = torch.ones(2) * 2
 53 |         targets = torch.ones(2)
 54 |         actual_score = self.metric(inputs, targets)
 55 |         expected_score = torch.sum(
 56 |             torch.exp((inputs - targets) / self.metric.pos_factor) - 1
 57 |         )
 58 |         self.assertEqual(expected_score, actual_score)
 59 | 
 60 | 
 61 | class TestRMSE(unittest.TestCase):
 62 |     def setUp(self):
 63 |         self.metric = metrics.RMSELoss()
 64 | 
 65 |     def test_update(self):
 66 |         expected_sse, batch_size = self._add_one_batch()
 67 |         self.assertEqual(1, self.metric.sample_counter)
 68 |         self.assertEqual(expected_sse, self.metric.losses[0])
 69 |         self.assertEqual(batch_size, self.metric.sizes[0])
 70 | 
 71 |     def test_reset(self):
 72 |         self._add_one_batch()
 73 |         self.metric.reset()
 74 |         self.assertEqual(0, self.metric.sample_counter)
 75 |         self.assertEqual(0, self.metric.losses.sum())
 76 |         self.assertEqual(0, self.metric.sizes.sum())
 77 | 
 78 |     def _add_one_batch(self):
 79 |         batch_size = 100
 80 |         inputs = torch.ones(batch_size) * 2
 81 |         targets = torch.zeros_like(inputs)
 82 |         summed_squares = torch.sum((inputs - targets) ** 2)
 83 | 
 84 |         self.metric.update(inputs, targets)
 85 | 
 86 |         return summed_squares, batch_size
 87 | 
 88 |     def test_compute(self):
 89 |         batch_sizes = [3000, 3000, 3000, 1000]
 90 |         inputs = torch.randn(sum(batch_sizes)) + 100
 91 |         targets = torch.randn_like(inputs)
 92 |         expected_rmse = torch.sqrt(torch.mean((inputs - targets) ** 2))
 93 | 
 94 |         batched_inputs = torch.split(inputs, batch_sizes)
 95 |         batched_targets = torch.split(targets, batch_sizes)
 96 |         for inp, tgt in zip(batched_inputs, batched_targets):
 97 |             self.metric.update(inp, tgt)
 98 |         actual_rmse = self.metric.compute()
 99 | 
100 |         self.assertAlmostEqual(expected_rmse.item(), actual_rmse.item(), delta=0.001)
101 | 
102 |     def test_compute_fails_on_empty_metric(self):
103 |         with self.assertRaises(RuntimeError):
104 |             self.metric.compute()
105 | 
106 | 
107 | class TestMeanMetric(unittest.TestCase):
108 |     def setUp(self):
109 |         self.mean_metric = metrics.SimpleMetric()
110 |         self.sum_metric = metrics.SimpleMetric(reduction="sum")
111 | 
112 |     def test_update(self):
113 |         expected_loss, batch_size = self._add_one_batch()
114 |         self.assertEqual(1, self.mean_metric.sample_counter)
115 |         self.assertEqual(expected_loss, self.mean_metric.losses[0])
116 |         self.assertEqual(batch_size, self.mean_metric.sizes[0])
117 | 
118 |     def test_reset(self):
119 |         self._add_one_batch()
120 |         self.mean_metric.reset()
121 |         self.assertEqual(0, self.mean_metric.sample_counter)
122 |         self.assertEqual(0, self.mean_metric.losses.sum())
123 |         self.assertEqual(0, self.mean_metric.sizes.sum())
124 | 
125 |     def _add_one_batch(self):
126 |         batch_size = 100
127 |         loss = torch.tensor(500)
128 | 
129 |         self.mean_metric.update(loss, batch_size)
130 | 
131 |         return loss, batch_size
132 | 
133 |     def test_compute_mean(self):
134 |         batch_sizes = [512] * 50 + [100]
135 |         losses = torch.randn(sum(batch_sizes)) + 2
136 |         expected_loss = losses.mean()
137 | 
138 |         batched_inputs = torch.split(losses, batch_sizes)
139 |         for inp, sizes in zip(batched_inputs, batch_sizes):
140 |             self.mean_metric.update(inp.mean(), sizes)
141 |         actual_loss = self.mean_metric.compute()
142 | 
143 |         self.assertAlmostEqual(expected_loss.item(), actual_loss.item(), places=5)
144 | 
145 |     def test_compute_sum(self):
146 |         batch_sizes = [512] * 50 + [100]
147 |         losses = torch.randn(sum(batch_sizes)) + 2
148 |         expected_loss = losses.sum()
149 | 
150 |         batched_inputs = torch.split(losses, batch_sizes)
151 |         for inp, sizes in zip(batched_inputs, batch_sizes):
152 |             self.sum_metric.update(inp.sum(), sizes)
153 |         actual_loss = self.sum_metric.compute()
154 | 
155 |         self.assertAlmostEqual(expected_loss.item(), actual_loss.item(), delta=0.1)
156 | 
157 |     def test_compute_fails_on_empty_metric(self):
158 |         with self.assertRaises(RuntimeError):
159 |             self.mean_metric.compute()
160 |         with self.assertRaises(RuntimeError):
161 |             self.sum_metric.compute()
162 | 


--------------------------------------------------------------------------------
/tests/lightning_tests/test_pretraining.py:
--------------------------------------------------------------------------------
  1 | import unittest
  2 | from unittest import mock
  3 | 
  4 | import torch
  5 | 
  6 | from lightning import pretraining
  7 | 
  8 | 
  9 | class TestUnsupervisedPretraining:
 10 |     def test_mode_hparam(self):
 11 |         self.assertLess(1, len(self.net.hparams))
 12 |         self.assertIn("mode", self.net.hparams)
 13 |         self.assertEqual("metric", self.net.hparams["mode"])
 14 | 
 15 |     @torch.no_grad()
 16 |     def test_encoder(self):
 17 |         inputs = torch.randn(16, 14, 30)
 18 |         outputs = self.net.encoder(inputs)
 19 |         self.assertEqual(self.encoder_shape, outputs.shape)
 20 | 
 21 |     @torch.no_grad()
 22 |     def test_get_anchor_query_embeddings(self):
 23 |         anchors = torch.randn(16, 14, 30)
 24 |         queries = torch.randn(16, 14, 30)
 25 |         anchor_embeddings, query_embeddings = self.net._get_anchor_query_embeddings(
 26 |             anchors, queries
 27 |         )
 28 | 
 29 |         self.assertEqual(self.encoder_shape, anchor_embeddings.shape)
 30 |         self.assertEqual(self.encoder_shape, query_embeddings.shape)
 31 | 
 32 |         for norm in torch.norm(anchor_embeddings, dim=1).tolist():
 33 |             self.assertAlmostEqual(1.0, norm, places=6)
 34 |         for norm in torch.norm(query_embeddings, dim=1).tolist():
 35 |             self.assertAlmostEqual(1.0, norm, places=6)
 36 | 
 37 |     @torch.no_grad()
 38 |     def test_forward(self):
 39 |         anchors = torch.randn(16, 14, 30)
 40 |         anchors[5] = 1.0
 41 |         queries = torch.randn(16, 14, 30)
 42 |         queries[5] = 1.0
 43 |         self.net.eval()
 44 |         distances = self.net(anchors, queries)
 45 | 
 46 |         self.assertEqual(torch.Size((16,)), distances.shape)
 47 |         self.assertAlmostEqual(0.0, distances[5].item(), delta=1e-6)
 48 |         self.assertNotAlmostEqual(0.0, distances.sum().item())
 49 | 
 50 |     def test_batch_independence(self):
 51 |         torch.autograd.set_detect_anomaly(True)
 52 | 
 53 |         anchors = torch.randn(16, 14, 30)
 54 |         queries = torch.randn(16, 14, 30)
 55 |         anchors.requires_grad = True
 56 | 
 57 |         # Compute forward pass in eval mode to deactivate batch norm
 58 |         self.net.eval()
 59 |         outputs = self.net(anchors, queries)
 60 |         self.net.train()
 61 | 
 62 |         # Mask loss for certain samples in batch
 63 |         batch_size = outputs.shape[0]
 64 |         mask_idx = torch.randint(0, batch_size, ())
 65 |         mask = torch.ones_like(outputs)
 66 |         mask[mask_idx] = 0
 67 |         output = outputs * mask
 68 | 
 69 |         # Compute backward pass
 70 |         loss = output.mean()
 71 |         loss.backward(retain_graph=True)
 72 | 
 73 |         # Check if gradient exists and is zero for masked samples
 74 |         for i, grad in enumerate(anchors.grad[:batch_size]):
 75 |             if i == mask_idx:
 76 |                 self.assertTrue(torch.all(grad == 0).item())
 77 |             else:
 78 |                 self.assertTrue(not torch.all(grad == 0))
 79 |         anchors.grad = None
 80 | 
 81 |         torch.autograd.set_detect_anomaly(False)
 82 | 
 83 |     def test_all_parameters_updated(self):
 84 |         optim = torch.optim.SGD(self.net.parameters(), lr=0.1)
 85 | 
 86 |         inputs = (
 87 |             torch.randn(16, 14, 30),
 88 |             torch.randn(16, 14, 30),
 89 |             torch.randn(16),
 90 |             torch.randn(16),
 91 |         )
 92 |         loss = self.net.training_step(inputs, batch_idx=0)
 93 |         loss.backward()
 94 |         optim.step()
 95 | 
 96 |         for param_name, param in self.net.named_parameters():
 97 |             if param.requires_grad:
 98 |                 with self.subTest(name=param_name):
 99 |                     self.assertIsNotNone(param.grad)
100 |                     self.assertNotEqual(0.0, torch.sum(param.grad ** 2))
101 | 
102 |     @mock.patch("pytorch_lightning.LightningModule.log")
103 |     @torch.no_grad()
104 |     def test_metric_val_reduction(self, mock_log):
105 |         paired_batches = 600
106 |         embedding_dim = self.encoder_shape[-1]
107 |         anchor_embeddings = torch.randn(paired_batches, 32, embedding_dim)
108 |         query_embeddings = torch.randn(paired_batches, 32, embedding_dim)
109 |         domain_predictions = torch.randn(paired_batches, 32)
110 |         self._mock_predictions(anchor_embeddings, query_embeddings, domain_predictions)
111 | 
112 |         self._feed_dummy_val(paired_batches)
113 | 
114 |         regression_loss = torch.pairwise_distance(
115 |             anchor_embeddings.view(-1, embedding_dim),
116 |             query_embeddings.view(-1, embedding_dim),
117 |             eps=1e-8,
118 |         )
119 |         regression_loss = torch.mean(regression_loss ** 2)
120 |         domain_loss = torch.tensor(0.0)
121 |         checkpoint_score = regression_loss - 0.1 * domain_loss
122 |         expected_logs = {
123 |             "val/regression_loss": (regression_loss, 0.0001),
124 |             "val/domain_loss": (domain_loss, 0.0001),
125 |             "val/checkpoint_score": (checkpoint_score, 0.0001),
126 |         }
127 |         self.net.validation_epoch_end([])
128 | 
129 |         self._assert_logs(mock_log, expected_logs)
130 | 
131 |     def test_metric_val_updates(self):
132 |         paired_batches = 600
133 |         embedding_dim = self.encoder_shape[-1]
134 |         anchor_embeddings = torch.randn(paired_batches, 32, embedding_dim)
135 |         query_embeddings = torch.randn(paired_batches, 32, embedding_dim)
136 |         domain_predictions = torch.randn(paired_batches, 32)
137 |         self._mock_predictions(anchor_embeddings, query_embeddings, domain_predictions)
138 | 
139 |         self._feed_dummy_val(paired_batches)
140 |         self.assertEqual(paired_batches, self.net.regression_metric.sample_counter)
141 |         self.assertEqual(paired_batches, self.net.domain_metric.sample_counter)
142 | 
143 |     def _mock_predictions(self, anchor_embeddings, query_embeddings, domain_predictions):
144 |         embeddings = torch.cat([anchor_embeddings, query_embeddings], dim=1)
145 |         self.net.encoder.forward = mock.MagicMock(side_effect=embeddings)
146 | 
147 |     def _feed_dummy_val(self, num_batches):
148 |         for i in range(num_batches):
149 |             self.net.validation_step(
150 |                 (
151 |                     torch.zeros(32, 14, 30),
152 |                     torch.zeros(32, 14, 30),
153 |                     torch.zeros(32),
154 |                     torch.zeros(32),
155 |                 ),
156 |                 batch_idx=i,
157 |                 dataloader_idx=0,
158 |             )
159 | 
160 |     def _assert_logs(self, mock_log, expected_logs):
161 |         mock_log.assert_called()
162 |         for call in mock_log.mock_calls:
163 |             metric = call[1][0]
164 |             self.assertIn(metric, expected_logs, "Unexpected logged metric found.")
165 |             expected_value, delta = expected_logs[metric]
166 |             expected_value = expected_value.item()
167 |             actual_value = call[1][1].item()
168 |             with self.subTest(metric):
169 |                 self.assertAlmostEqual(expected_value, actual_value, delta=delta)
170 | 
171 | 
172 | class TestCnnUnsupervisedPretraining(unittest.TestCase, TestUnsupervisedPretraining):
173 |     def setUp(self):
174 |         self.encoder_shape = torch.Size((16, 64))
175 |         self.net = pretraining.UnsupervisedPretraining(
176 |             in_channels=14,
177 |             seq_len=30,
178 |             num_layers=4,
179 |             kernel_size=3,
180 |             base_filters=16,
181 |             latent_dim=64,
182 |             dropout=0.1,
183 |             domain_tradeoff=0.0,
184 |             domain_disc_dim=32,
185 |             num_disc_layers=2,
186 |             weight_decay=0,
187 |             lr=0.01,
188 |         )
189 | 


--------------------------------------------------------------------------------
/tests/model_tests/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/tilman151/self-supervised-ssl/745d4b8b06a946f1be1e2c8bc15e4b0b55c7c4ff/tests/model_tests/__init__.py


--------------------------------------------------------------------------------
/tests/model_tests/test_layers.py:
--------------------------------------------------------------------------------
 1 | import unittest
 2 | 
 3 | import torch
 4 | 
 5 | from models import layers
 6 | 
 7 | 
 8 | class TestDeflatten(unittest.TestCase):
 9 |     def test_deflattening(self):
10 |         inputs = torch.randn(16, 512)
11 |         expected_shape = (16, 8, 64)
12 |         module = layers.DeFlatten(64, 8)
13 | 
14 |         deflattened = module(inputs)
15 |         reflattened = deflattened.view(16, -1)
16 | 
17 |         self.assertEqual(expected_shape, deflattened.shape)
18 |         self.assertEqual(0.0, torch.sum(inputs - reflattened).item())
19 | 
20 | 
21 | class TestGradientReversal(unittest.TestCase):
22 |     def test_against_identity(self):
23 |         inputs = torch.randn(16, 64, 1)
24 |         weight = torch.randn(1, 32, 64)
25 |         weight.requires_grad = True
26 |         grad_reversal = layers.GradientReversalLayer()
27 | 
28 |         normal_loss = torch.mean(weight @ inputs)
29 |         normal_loss.backward()
30 |         normal_gradient = weight.grad.clone()
31 |         weight.grad = None
32 | 
33 |         reverse_loss = torch.mean(grad_reversal(weight @ inputs))
34 |         reverse_loss.backward()
35 |         reverse_gradient = weight.grad.clone()
36 | 
37 |         self.assertEqual(0.0, torch.sum(normal_gradient + reverse_gradient))
38 | 


--------------------------------------------------------------------------------
/tests/model_tests/test_networks.py:
--------------------------------------------------------------------------------
 1 | import unittest
 2 | 
 3 | import torch
 4 | 
 5 | from models import networks
 6 | 
 7 | 
 8 | class ModelTestsMixin:
 9 |     @torch.no_grad()
10 |     def test_shape(self):
11 |         outputs = self.net(self.test_inputs)
12 |         self.assertEqual(self.expected_shape, outputs.shape)
13 | 
14 |     @torch.no_grad()
15 |     @unittest.skipUnless(torch.cuda.is_available(), "No GPU was detected")
16 |     def test_device_moving(self):
17 |         self.net.eval()
18 | 
19 |         torch.manual_seed(42)
20 |         outputs_cpu = self.net(self.test_inputs)
21 | 
22 |         net_on_gpu = self.net.to("cuda:0")
23 |         torch.manual_seed(42)
24 |         outputs_gpu = net_on_gpu(self.test_inputs.to("cuda:0"))
25 | 
26 |         torch.manual_seed(42)
27 |         net_back_on_cpu = net_on_gpu.cpu()
28 |         outputs_back_on_cpu = net_back_on_cpu(self.test_inputs)
29 | 
30 |         self.assertAlmostEqual(
31 |             0.0, torch.sum(outputs_cpu - outputs_gpu.cpu()).item(), delta=10e-6
32 |         )
33 |         self.assertAlmostEqual(
34 |             0.0, torch.sum(outputs_cpu - outputs_back_on_cpu).item(), delta=10e-6
35 |         )
36 | 
37 |     def test_batch_independence(self):
38 |         inputs = self.test_inputs.clone()
39 |         inputs.requires_grad = True
40 | 
41 |         # Compute forward pass in eval mode to deactivate batch norm
42 |         self.net.eval()
43 |         outputs = self.net(inputs)
44 |         self.net.train()
45 | 
46 |         # Mask loss for certain samples in batch
47 |         batch_size = inputs.shape[0]
48 |         mask_idx = torch.randint(0, batch_size, ())
49 |         mask = torch.ones_like(outputs)
50 |         mask[mask_idx] = 0
51 |         outputs = outputs * mask
52 | 
53 |         # Compute backward pass
54 |         loss = outputs.mean()
55 |         loss.backward()
56 | 
57 |         # Check if gradient exists and is zero for masked samples
58 |         for i, grad in enumerate(inputs.grad):
59 |             if i == mask_idx:
60 |                 self.assertTrue(torch.all(grad == 0).item())
61 |             else:
62 |                 self.assertTrue(not torch.all(grad == 0))
63 | 
64 |     def test_all_parameters_updated(self):
65 |         optim = torch.optim.SGD(self.net.parameters(), lr=0.1)
66 | 
67 |         outputs = self.net(self.test_inputs)
68 |         loss = outputs.mean()
69 |         loss.backward()
70 |         optim.step()
71 | 
72 |         for param_name, param in self.net.named_parameters():
73 |             if param.requires_grad:
74 |                 with self.subTest(name=param_name):
75 |                     self.assertIsNotNone(param.grad)
76 |                     self.assertNotEqual(0.0, torch.sum(param.grad ** 2))
77 | 
78 | 
79 | class TestEncoder(unittest.TestCase, ModelTestsMixin):
80 |     def setUp(self):
81 |         self.net = networks.Encoder(14, 16, 3, 6, 64, 30, 0.1, True)
82 |         self.test_inputs = torch.randn(16, 14, 30)
83 |         self.expected_shape = (16, 64)
84 | 
85 | 
86 | class TestDecoder(unittest.TestCase, ModelTestsMixin):
87 |     def setUp(self):
88 |         self.net = networks.Decoder(14, 16, 3, 6, 64, 30, 0.1)
89 |         self.test_inputs = torch.randn(16, 64)
90 |         self.expected_shape = (16, 14, 30)
91 | 
92 | 
93 | class TestRegressor(unittest.TestCase, ModelTestsMixin):
94 |     def setUp(self):
95 |         self.net = networks.Regressor(64)
96 |         self.test_inputs = torch.randn(16, 64)
97 |         self.expected_shape = (16,)
98 | 


--------------------------------------------------------------------------------
/tests/test_building.py:
--------------------------------------------------------------------------------
  1 | import unittest
  2 | from unittest import mock
  3 | 
  4 | from building import build, get_logdir
  5 | from lightning import loggers
  6 | 
  7 | 
  8 | class TestBuildingFunctions(unittest.TestCase):
  9 |     def setUp(self):
 10 |         self.config = {
 11 |             "num_layers": 8,
 12 |             "base_filters": 16,
 13 |             "domain_tradeoff": 1.0,
 14 |             "latent_dim": 32,
 15 |             "dropout": 0.1,
 16 |             "num_disc_layers": 1,
 17 |             "lr": 0.01,
 18 |             "batch_size": 512,
 19 |         }
 20 |         self.pretraining_config = {
 21 |             "domain_tradeoff": 2.0,
 22 |             "dropout": 0.3,
 23 |             "lr": 0.02,
 24 |             "batch_size": 256,
 25 |             "distance_mode": "piecewise",
 26 |         }
 27 | 
 28 |     @mock.patch("lightning.loggers.MinEpochModelCheckpoint")
 29 |     @mock.patch("lightning.loggers.MLTBLogger")
 30 |     @mock.patch("building.build.build_pretraining_from_config")
 31 |     @mock.patch("building.build.build_datamodule")
 32 |     @mock.patch("building.build.build_trainer")
 33 |     def test_build_pretraining_metric(
 34 |         self,
 35 |         mock_build_trainer,
 36 |         mock_datamodule,
 37 |         mock_pretraining_from_config,
 38 |         mock_logger,
 39 |         mock_checkpoint,
 40 |     ):
 41 |         self._check_build_pretraining(
 42 |             mock_build_trainer,
 43 |             mock_checkpoint,
 44 |             mock_datamodule,
 45 |             mock_logger,
 46 |             mock_pretraining_from_config,
 47 |             "metric",
 48 |         )
 49 | 
 50 |     @mock.patch("lightning.loggers.MinEpochModelCheckpoint")
 51 |     @mock.patch("lightning.loggers.MLTBLogger")
 52 |     @mock.patch("building.build.build_autoencoder_from_config")
 53 |     @mock.patch("building.build.build_datamodule")
 54 |     @mock.patch("building.build.build_trainer")
 55 |     def test_build_pretraining_autoencoder(
 56 |         self,
 57 |         mock_build_trainer,
 58 |         mock_datamodule,
 59 |         mock_pretraining_from_config,
 60 |         mock_logger,
 61 |         mock_checkpoint,
 62 |     ):
 63 |         self._check_build_pretraining(
 64 |             mock_build_trainer,
 65 |             mock_checkpoint,
 66 |             mock_datamodule,
 67 |             mock_logger,
 68 |             mock_pretraining_from_config,
 69 |             "autoencoder",
 70 |         )
 71 | 
 72 |     def _check_build_pretraining(
 73 |         self,
 74 |         mock_build_trainer,
 75 |         mock_checkpoint,
 76 |         mock_datamodule,
 77 |         mock_logger,
 78 |         mock_pretraining_from_config,
 79 |         mode,
 80 |     ):
 81 |         mock_logger.return_value = "mock_logger"
 82 |         mock_checkpoint.return_value = "mock_checkpoint"
 83 |         build.build_pretraining(
 84 |             2,
 85 |             1,
 86 |             0.8,
 87 |             0.5,
 88 |             self.config,
 89 |             self.pretraining_config,
 90 |             "cnn",
 91 |             mode,
 92 |             False,
 93 |             1,
 94 |             42,
 95 |             "version",
 96 |         )
 97 |         mock_logger.assert_called_with(
 98 |             get_logdir(),
 99 |             loggers.pretraining_experiment_name(2, 1),
100 |             tag="version",
101 |             tensorboard_struct={
102 |                 "pb": 0.8,
103 |                 "dt": self.pretraining_config["domain_tradeoff"],
104 |             },
105 |         )
106 |         mock_build_trainer.assert_called_with(
107 |             "mock_logger",
108 |             "mock_checkpoint",
109 |             max_epochs=100,
110 |             val_interval=1.0,
111 |             gpu=1,
112 |             seed=42,
113 |         )
114 |         mock_datamodule.assert_called_with(
115 |             2,
116 |             1,
117 |             0.8,
118 |             0.5,
119 |             self.pretraining_config["batch_size"],
120 |             True,
121 |             "piecewise",
122 |             1,
123 |         )
124 |         mock_pretraining_from_config.assert_called_with(
125 |             self.config,
126 |             self.pretraining_config,
127 |             mock_datamodule().window_size,
128 |             "cnn",
129 |             False,
130 |             True,
131 |         )
132 | 
133 |     @mock.patch("lightning.pretraining.UnsupervisedPretraining")
134 |     def test_build_pretraining_from_config(self, mock_pretraining):
135 |         self._check_pretraining_from_config(
136 |             build.build_pretraining_from_config, mock_pretraining
137 |         )
138 | 
139 |     @mock.patch("lightning.autoencoder.AutoencoderPretraining")
140 |     def test_build_autoencoder_from_config(self, mock_autoencoder):
141 |         self._check_pretraining_from_config(
142 |             build.build_autoencoder_from_config, mock_autoencoder
143 |         )
144 | 
145 |     def _check_pretraining_from_config(self, building_func, mock_pretraining):
146 |         building_func(
147 |             self.config, self.pretraining_config, 30, "cnn", False, use_adaption=True
148 |         )
149 |         mock_pretraining.assert_called_with(
150 |             in_channels=14,
151 |             seq_len=30,
152 |             num_layers=self.config["num_layers"],
153 |             kernel_size=3,
154 |             base_filters=self.config["base_filters"],
155 |             latent_dim=self.config["latent_dim"],
156 |             encoder="cnn",
157 |             dropout=self.pretraining_config["dropout"],
158 |             domain_tradeoff=self.pretraining_config["domain_tradeoff"],
159 |             domain_disc_dim=self.config["latent_dim"],
160 |             num_disc_layers=self.config["num_disc_layers"],
161 |             lr=self.pretraining_config["lr"],
162 |             weight_decay=0.0,
163 |             record_embeddings=False,
164 |         )
165 |         building_func(
166 |             self.config, self.pretraining_config, 30, "cnn", False, use_adaption=False
167 |         )
168 |         mock_pretraining.assert_called_with(
169 |             in_channels=14,
170 |             seq_len=30,
171 |             num_layers=self.config["num_layers"],
172 |             kernel_size=3,
173 |             base_filters=self.config["base_filters"],
174 |             latent_dim=self.config["latent_dim"],
175 |             encoder="cnn",
176 |             dropout=self.pretraining_config["dropout"],
177 |             domain_tradeoff=0.0,
178 |             domain_disc_dim=self.config["latent_dim"],
179 |             num_disc_layers=self.config["num_disc_layers"],
180 |             lr=self.pretraining_config["lr"],
181 |             weight_decay=0.0,
182 |             record_embeddings=False,
183 |         )
184 | 
185 |     @mock.patch("pytorch_lightning.callbacks.ModelCheckpoint")
186 |     @mock.patch("lightning.loggers.MLTBLogger")
187 |     @mock.patch("building.build.build_baseline_from_config")
188 |     @mock.patch("datasets.BaselineDataModule")
189 |     @mock.patch("building.build.build_trainer")
190 |     def test_build_baseline(
191 |         self,
192 |         mock_build_trainer,
193 |         mock_datamodule,
194 |         mock_baseline_from_config,
195 |         mock_logger,
196 |         mock_checkpoint,
197 |     ):
198 |         mock_logger.return_value = "mock_logger"
199 |         mock_checkpoint.return_value = "mock_checkpoint"
200 |         with self.subTest("no_encoder"):
201 |             build.build_baseline(2, 1.0, self.config, "cnn", None, 1, 42, "version")
202 |             mock_logger.assert_called_with(
203 |                 get_logdir(),
204 |                 loggers.baseline_experiment_name(2),
205 |                 tag="version",
206 |             )
207 |             mock_build_trainer.assert_called_with(
208 |                 "mock_logger",
209 |                 "mock_checkpoint",
210 |                 max_epochs=100,
211 |                 val_interval=1.0,
212 |                 gpu=1,
213 |                 seed=42,
214 |             )
215 |             mock_datamodule.assert_called_with(
216 |                 fd_source=2, percent_fail_runs=1.0, batch_size=self.config["batch_size"]
217 |             )
218 |             mock_baseline_from_config.assert_called_with(
219 |                 self.config,
220 |                 mock_datamodule().window_size,
221 |                 "cnn",
222 |                 None,
223 |                 False,
224 |             )
225 |         with self.subTest("with_encoder"):
226 |             build.build_baseline(
227 |                 2, 1.0, self.config, "cnn", "encoder_path", 1, 42, "version"
228 |             )
229 |             mock_build_trainer.assert_called_with(
230 |                 "mock_logger",
231 |                 "mock_checkpoint",
232 |                 max_epochs=100,
233 |                 val_interval=1.0,
234 |                 gpu=1,
235 |                 seed=42,
236 |             )
237 |             mock_datamodule.assert_called_with(
238 |                 fd_source=2, percent_fail_runs=1.0, batch_size=self.config["batch_size"]
239 |             )
240 |             mock_baseline_from_config.assert_called_with(
241 |                 self.config,
242 |                 mock_datamodule().window_size,
243 |                 "cnn",
244 |                 "encoder_path",
245 |                 False,
246 |             )
247 | 
248 |     @mock.patch("lightning.baseline.Baseline")
249 |     def test_build_baseline_from_config(self, mock_baseline):
250 |         with self.subTest("no_encoder"):
251 |             build.build_baseline_from_config(self.config, 30, "cnn", None)
252 |             self._assert_baseline_build_correctly(mock_baseline)
253 |             mock_baseline().load_encoder.assert_not_called()
254 |         with self.subTest("with_encoder"):
255 |             build.build_baseline_from_config(self.config, 30, "cnn", "encoder_path")
256 |             self._assert_baseline_build_correctly(mock_baseline)
257 |             mock_baseline().load_encoder.assert_called_with("encoder_path")
258 | 
259 |     def _assert_baseline_build_correctly(self, mock_baseline):
260 |         mock_baseline.assert_called_with(
261 |             in_channels=14,
262 |             seq_len=30,
263 |             num_layers=self.config["num_layers"],
264 |             kernel_size=3,
265 |             base_filters=self.config["base_filters"],
266 |             latent_dim=self.config["latent_dim"],
267 |             dropout=self.config["dropout"],
268 |             optim_type="adam",
269 |             lr=self.config["lr"],
270 |             record_embeddings=False,
271 |             encoder="cnn",
272 |         )
273 | 


--------------------------------------------------------------------------------