MIDI-92: Improvements to Discriminator and Data Normalization (#2)

* experiment with noise size

* combine all data for loader

* batch_size

* broken TGAN

* changed to 2 epochs

* code from vqvae-midi

* encode, decode function

* copied transformer code to revise

* rename to vqvae

* refactor 01

* vqvae training for 1 chanel

* cleaned dataset

* interpolation eval

* remove tanh

* compute statistics for all data

* normalized data, tanh

* denormalization in evals

* Add ResidualBlocks

* update config

* Update README

* preprocessing ltafdb

* change name

Author: SamuelJanas

README.md

@@ -1,6 +1,6 @@
 # ecg-gan
 
-Check out this [wandb run](https://wandb.ai/sjanas/ECG%20GAN/runs/bpr8skng) to get a feel for the baseline model or run `python train.py` to reproduce the results
+Check out this [wandb run](https://wandb.ai/sjanas/ECG%20GAN/runs/fz2ptaxy) to get a feel for the baseline model or run `python train.py` to reproduce the results
 
 ### Suggested reading
 

configs/config.yaml

@@ -1,13 +1,13 @@
 train:
-  batch_size: 64
-  epochs: 5
+  batch_size: 16
+  epochs: 3
   log_interval: 100
-  generator_lr: 1e-4
+  generator_lr: 2e-4
   discriminator_adam_lr: 1e-5
-  num_workers: 4
+  num_workers: 6
   # For SGD optimizer
   use_sgd: true
-  discriminator_sgd_lr: 2e-4 # It's slightly slower and needs higher learning rate
+  discriminator_sgd_lr: 1e-4 # It's slightly slower and needs higher learning rate
   # for continuing training
   load_checkpoint: # change from None to checkpoint path to continue training
   more_epochs: 20 # How many more epochs to train
@@ -29,9 +29,9 @@ discriminator:
   beta: 0.5
 
 generator:
-  noise_size: 200
+  noise_size: 300
   beta: 0.5
 
 project: "ECG GAN"
 run_date: ${now:%Y_%m_%d_%H_%M}
-run_name: "ECG_GAN_${run_date}"
+run_name: "residual_discriminator_${run_date}"

configs/config_tgan.yaml

@@ -0,0 +1,39 @@
+train:
+  batch_size: 16
+  epochs: 5
+  log_interval: 100
+  generator_lr: 1e-4
+  discriminator_lr: 1e-5
+  num_workers: 4
+
+  # for continuing training
+  load_checkpoint: # change from None to checkpoint path to continue training
+  more_epochs: 20 # How many more epochs to train
+
+logger:
+  checkpoint_path: "checkpoints/"
+  chart_path: "tmp/"
+
+system:
+  device: "cuda:0"
+  seed: 23
+
+data:
+  channels: 1
+  size: 1000
+
+discriminator:
+  n_layers: 1
+  n_channel: 1
+  kernel_size: 8
+  dropout: 0.0
+
+generator:
+  noise_size: 100
+  beta: 0.5
+  n_layers: 1
+  hidden_dim: 256
+
+project: "ECG TGAN"
+run_date: ${now:%Y_%m_%d_%H_%M}
+run_name: "ECG_TGAN_${run_date}"

configs/config_transformer.yaml

@@ -0,0 +1,24 @@
+latent_dim: 256
+image_size: 256
+num_codebook_vectors: 1024
+beta: 0.25
+image_channels: 1  # Since 1D data?
+dataset_path: "./data"
+checkpoint_path: "./checkpoints/last_ckpt.pt"
+batch_size: 20
+epochs: 100
+learning_rate: 2.25e-05
+beta1: 0.5
+beta2: 0.9
+disc_start: 10000
+disc_factor: 1.0
+l2_loss_factor: 1.0
+perceptual_loss_factor: 1.0
+pkeep: 0.5
+sos_token: 0
+
+system:
+  device: 'cuda:0'
+
+# Parameters specific to our data shape
+seq_len: 1000

configs/config_vqvae.yaml

@@ -0,0 +1,45 @@
+train:
+  batch_size: 16
+  epochs: 2
+  log_interval: 100
+  lr: 2e-4
+  num_workers: 4
+
+logger:
+  checkpoint_path: "checkpoints/"
+  chart_path: "tmp/"
+
+system:
+  device: "cuda:0"
+  seed: 23
+
+data:
+  channels: 1
+  size: 1000
+
+model:
+  # Structure/Architecture
+  output_features_filters: 1
+  augment_output_features: False
+  input_features_filters: 1  # Aligned with the number of channels
+  augment_input_features: False
+  output_features_dim: 1000  # Features size post decoding
+  input_features_dim: 1000  # Aligned with the number of data points in each channel
+  input_features_type: 'mfcc'
+  # Encoder-Decoder details
+  num_hiddens: 128
+  num_residual_layers: 2
+  num_residual_hiddens: 16
+  use_kaiming_normal: True
+  # VQ details
+  embedding_dim: 100 # We might want to lower this value in the future, size of embedding vectors
+  num_embeddings: 64 # K, needs empirical tuning
+  commitment_cost: 0.25
+  decay: 0.99
+  # Misc
+  record_codebook_stats: False
+  verbose: False
+
+project: "ECG GAN"
+run_date: ${now:%Y_%m_%d_%H_%M}
+run_name: "ECG_GAN_${run_date}"

data_showcase.ipynb

@@ -0,0 +1,14 @@
+import torch
+from omegaconf import OmegaConf
+from huggingface_hub.file_download import hf_hub_download
+
+
+def load_checkpoint(ckpt_path: str = None, omegaconf: bool = True):
+    if ckpt_path is not None:
+        checkpoint = torch.load(ckpt_path)
+    else:
+        checkpoint = torch.load(hf_hub_download("SneakyInsect/GANs", filename="double_generator_update_2023_09_26_21_23_3.pt"))
+    cfg = checkpoint["config"]
+    if omegaconf:
+        cfg = OmegaConf.create(cfg)
+    return checkpoint, cfg

evals/checkpoint_utils.py

@@ -0,0 +1,79 @@
+import torch
+import imageio
+import numpy as np
+from matplotlib import pyplot as plt
+
+from model.dcgan import Generator
+from utils.dataloaders import set_seed
+from evals.checkpoint_utils import load_checkpoint
+
+
+def interpolate(noise1: torch.Tensor, noise2: torch.Tensor, num_interpolations: int) -> torch.Tensor:
+    """
+    Create interpolated noises between noise1 and noise2, including noise1 at the start and noise2 at the end.
+    """
+    alphas = torch.linspace(0, 1, num_interpolations + 2).to(noise1.device)  # +2 to account for noise1 and noise2
+    noises = [(1 - alpha) * noise1 + alpha * noise2 for alpha in alphas]
+    return torch.stack(noises, dim=0)
+
+
+@torch.no_grad()
+def save_interpolations(generator, noise1: torch.Tensor, noise2: torch.Tensor, num_interpolations: int, save_path: str):
+    generator.eval()
+
+    # create num_interpolation noises between noise1 and noise2
+    noises = interpolate(noise1, noise2, num_interpolations)
+
+    # generate fake data using the generator and noises
+    fake_data = generator(noises).detach().cpu().numpy()
+
+    # denormalize data from [-1, 1]
+    global_max = [9.494109153747559, 7.599456787109375]
+    global_min = [-10.515237808227539, -7.820725917816162]
+    fake_data = (fake_data + 1) * (global_max[0] - global_min[0]) / 2 + global_min[0]
+
+    # Visualize and save as gif
+    images = []
+    for i in range(len(fake_data)):
+        fig, ax = plt.subplots(figsize=[10, 4], gridspec_kw={"hspace": 0})
+        ax.plot(fake_data[i, 0, :])  # Plotting the data
+        ax.set_title("Generated Data Visualization")
+
+        # Convert Figure to image and append to images list
+        fig.canvas.draw()
+        image = np.frombuffer(fig.canvas.tostring_rgb(), dtype="uint8")
+        image = image.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+        images.append(image)
+        plt.close(fig)
+
+    # Save images as gif
+    imageio.mimsave(save_path, images, duration=250)
+    print(f"Saved interpolation gif to {save_path}")
+
+
+if __name__ == "__main__":
+    checkpoint_path = (
+        "checkpoints/ECG_GAN_2023_09_29_10_00_1.pt"  # change this to string path of checkpoint if you want to load from local
+    )
+    number_of_interpolations = 20
+    save_path = "tmp/interpolation.gif"
+
+    # set seed for reproducibility
+    set_seed(23)
+
+    # load checkpoint, config
+    checkpoint, cfg = load_checkpoint(checkpoint_path)
+
+    # Initialize generator
+    generator = Generator(
+        noise_size=cfg.generator.noise_size,
+        output_size=cfg.data.size,
+    ).to(cfg.system.device)
+
+    generator.load_state_dict(checkpoint["generator_state_dict"])
+
+    # create 2 noise vectors
+    noise1 = torch.randn(cfg.generator.noise_size, cfg.data.channels, device=cfg.system.device)
+    noise2 = torch.randn(cfg.generator.noise_size, cfg.data.channels, device=cfg.system.device)
+
+    save_interpolations(generator, noise1, noise2, number_of_interpolations, save_path)

evals/interpolate_ecg.py

@@ -10,28 +10,46 @@ def weights_init(m):
         nn.init.constant_(m.bias.data, 0)
 
 
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(ResidualBlock, self).__init__()
+
+        self.residual = nn.Sequential(
+            nn.Conv1d(in_channels, out_channels, 4, 2, 1, bias=False),
+            nn.BatchNorm1d(out_channels),
+            nn.LeakyReLU(0.2, inplace=True),
+            # This line calculates the padding to ensure the size remains consistent
+            nn.ConstantPad1d((0, 1), 0),
+            nn.Conv1d(out_channels, out_channels, 4, 1, 1, bias=False),
+            nn.BatchNorm1d(out_channels),
+        )
+
+        self.shortcut = nn.Sequential()
+        if in_channels != out_channels:
+            self.shortcut.add_module("conv_shortcut", nn.Conv1d(in_channels, out_channels, 4, 2, 1, bias=False))
+            self.shortcut.add_module("bn_shortcut", nn.BatchNorm1d(out_channels))
+
+    def forward(self, x):
+        return nn.LeakyReLU(0.2, inplace=True)(self.residual(x) + self.shortcut(x))
+
+
 class Discriminator(nn.Module):
     def __init__(self, input_channels, input_size, neurons):
         super().__init__()
+
         layers = []
         prev_channels = input_channels
         for n in neurons:
-            layers.extend(
-                [
-                    nn.Conv1d(prev_channels, n, 4, 2, 1, bias=False),
-                    nn.BatchNorm1d(n) if prev_channels != input_channels else nn.Identity(),
-                    nn.LeakyReLU(0.2, inplace=True),
-                ]
-            )
+            layers.append(ResidualBlock(prev_channels, n))
             prev_channels = n
+
         layers.append(nn.Conv1d(prev_channels, 1, input_size // (2 ** len(neurons)), 1, 0, bias=False))
         layers.append(nn.Sigmoid())
 
         self.main = nn.Sequential(*layers)
 
     def forward(self, x):
-        x = self.main(x)
-        return x
+        return self.main(x)
 
 
 class Generator(nn.Module):