DualHeadSSIM

Dual-Head SSIM model. This is the improved SSIM model with two GRU encoders.

YiFan Zhang and Peter Thorburn. A Dual-Head Attention Model for Time Series Data Imputation, Volume 189, 2021, 106377, ISSN 0168-1699, Computers and Electronics in Agriculture.

Code Structure

/checkpoints ------- store trained model

Level3_1012 means this model imputes water level data with gap size 3, and it is tested on data from Oct to Dec (Detail of train/test data split is in the above paper)

/comparative ------- KNN.py stores all the implementation for KNN, EM and other traditional methods

/data ------- data set used for model train/test

OLD_nomiss.csv is the one used for the training model in the paper. To train the imputation model, we manually masked the available data as the 'missing data'. In this way, we can have groundtruth data to supervise training.

This csv is generated from the original WQ monitoring datasets. The data is published on Kaggle as Real Time Water Quality Data

/models ------- Dual-SSIM model: left/right encoder, decoder, cross attention

/paper_related ------------ Mainly including some plotting functions for all the figures listed in the paper

/results ------------ some imputation resutls generated by models.

Dual_SSIM_level_1012_ori means these imputation results are generated by a Dual-Head SSIM model. This model is trained on imputing water level data, and it is tested on data from Oct to Dec. ori means this is the results in the original scale. (Otherwise, scal means the results are normalized by MaxMinScaler)

/utils --------------- including some useful functions to help training. For example, the implementation of dilate loss. It also has the data processing functions to generate train/test samples for deep learning models such as MRNN and BRITS.

/mean.py ------------- The main functions to train/test the Dual-Head SSIM model.

An end-to-end example

The code is currently under development.

To train Dual-Head SSIM model on your own data:

1. Data preprocessing

   Get data and have a data preprocessing function similar to the one in /utils/prepare_QLD.py

2. In main.py

   1. Setup hyperparameters from line 230 to line 242
   2. Replace preprocessing data function in line 246
   3. Change line 250 and 251, based on your data preprocessing setting (in paper, the value p, k and q in Figure 1)
   4. Line 310, setup the checkpoints
   5. Line 318 to 361 are the codes for model training.
   6. Line 372 to 383 are the codes for model evaluation.
   7. Line 387 to 396 can output and save imputations.
   8. Line 400 to 450 can plot every imputation through the test data.

If you train on WQ data provided in this repo, just setup prepare_QLD.py (setup which variable as the target, p,k and q), check steps 2, 4, 5, and uncomment step 6.

The model train/evaluation for BRITS and MRNN can be found in repo BRITS_TEST