add lightning code, finetuning whisper, recommender system neural collaborative filtering

ML/Pytorch/more_advanced/VAE/lightning_vae/dataset.py

@@ -0,0 +1,60 @@
+# Imports
+import torch
+import torchvision.datasets as datasets  # Standard datasets
+import torchvision.transforms as transforms  # Transformations we can perform on our dataset for augmentation
+from torch.utils.data import DataLoader
+import pytorch_lightning as pl
+
+
+class MNISTDataModule(pl.LightningDataModule):
+    def __init__(self, batch_size, num_workers):
+        super().__init__()
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+
+    def setup(self, stage):
+        mnist_full = train_dataset = datasets.MNIST(
+            root="dataset/", train=True, transform=transforms.ToTensor(), download=True
+        )
+        self.mnist_test = datasets.MNIST(
+            root="dataset/", train=False, transform=transforms.ToTensor(), download=True
+        )
+        self.mnist_train, self.mnist_val = torch.utils.data.random_split(
+            mnist_full, [55000, 5000]
+        )
+
+    def train_dataloader(self):
+        return DataLoader(
+            self.mnist_train,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            persistent_workers=True,
+            shuffle=True,
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            self.mnist_val,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            persistent_workers=True,
+            shuffle=False,
+        )
+
+    def test_dataloader(self):
+        return DataLoader(
+            self.mnist_test,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            persistent_workers=True,
+            shuffle=False,
+        )
+
+
+# check that it works
+if __name__ == "__main__":
+    dm = MNISTDataModule()
+    dm.setup("fit")
+    print(len(dm.mnist_train))
+    print(len(dm.mnist_val))
+    print(len(dm.mnist_test))

ML/Pytorch/more_advanced/VAE/lightning_vae/model.py

@@ -0,0 +1,92 @@
+import torch
+import torchvision
+from torch import nn
+import pytorch_lightning as pl
+
+
+class VAEpl(pl.LightningModule):
+    def __init__(self, lr, input_dim=784, h_dim=200, z_dim=20):
+        super().__init__()
+        self.lr = lr
+        self.loss_fn = nn.BCELoss(reduction="sum")
+        self.input_dim = input_dim
+
+        # encoder
+        self.img_2hid = nn.Linear(input_dim, h_dim)
+        self.hid_2mu = nn.Linear(h_dim, z_dim)
+        self.hid_2sigma = nn.Linear(h_dim, z_dim)
+
+        # decoder
+        self.z_2hid = nn.Linear(z_dim, h_dim)
+        self.hid_2img = nn.Linear(h_dim, input_dim)
+        self.relu = nn.ReLU()
+        self.sigmoid = nn.Sigmoid()
+
+    def encode(self, x):
+        h = self.relu(self.img_2hid(x))
+        mu, sigma = self.hid_2mu(h), self.hid_2sigma(h)
+        return mu, sigma
+
+    def decode(self, z):
+        h = self.relu(self.z_2hid(z))
+        return torch.sigmoid(self.hid_2img(h))
+
+    def forward(self, x):
+        mu, sigma = self.encode(x)
+        epsilon = torch.randn_like(sigma)
+        z_new = mu + sigma * epsilon
+        x_reconstructed = self.decode(z_new)
+        return x_reconstructed, mu, sigma
+
+    def training_step(self, batch, batch_idx):
+        x, _ = batch
+        x = x.view(-1, self.input_dim)
+        x_reconstructed, mu, sigma = self.forward(x)
+        reconstruction_loss = self.loss_fn(x_reconstructed, x)
+        kl_div = -torch.sum(1 + torch.log(sigma.pow(2)) - mu.pow(2) - sigma.pow(2))
+        loss = reconstruction_loss + kl_div
+        self.log("train_loss", loss, sync_dist=True)
+
+        # add logging of images to tensorboard, x_reconstructed and x, so that
+        # it updates every step and we can the progress pictures in tensorboard
+        if batch_idx % 100 == 0:
+            # take out the first 8
+            x = x[:8]
+            x_reconstructed = x_reconstructed[:8]
+            grid = torchvision.utils.make_grid(x_reconstructed.view(-1, 1, 28, 28))
+            self.logger.experiment.add_image("reconstructed", grid, self.global_step)
+            grid = torchvision.utils.make_grid(x.view(-1, 1, 28, 28))
+            self.logger.experiment.add_image("original", grid, self.global_step)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        x, _ = batch
+        x = x.view(-1, self.input_dim)
+        x_reconstructed, mu, sigma = self.forward(x)
+        reconstruction_loss = self.loss_fn(x_reconstructed, x)
+        kl_div = -torch.sum(1 + torch.log(sigma.pow(2)) - mu.pow(2) - sigma.pow(2))
+        loss = reconstruction_loss + kl_div
+        self.log("val_loss", loss, sync_dist=True)
+        return loss
+
+    def test_step(self, batch, batch_idx):
+        x, _ = batch
+        x = x.view(-1, self.input_dim)
+        x_reconstructed, mu, sigma = self.forward(x)
+        reconstruction_loss = self.loss_fn(x_reconstructed, x)
+        kl_div = -torch.sum(1 + torch.log(sigma.pow(2)) - mu.pow(2) - sigma.pow(2))
+        loss = reconstruction_loss + kl_div
+        self.log("test_loss", loss, sync_dist=True)
+        return loss
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
+        return optimizer
+
+
+if __name__ == "__main__":
+    batch_size = 8
+    x = torch.randn(batch_size, 28 * 28 * 1)
+    vae_pl = VAEpl()
+    x_reconstructed, mu, sigma = vae_pl(x)
+    print(x_reconstructed.shape)

ML/Pytorch/more_advanced/VAE/lightning_vae/train.py

@@ -0,0 +1,49 @@
+import torch
+import torchvision.datasets as datasets  # Standard datasets
+from tqdm import tqdm
+from torch import nn, optim
+from torchvision import transforms
+from torchvision.utils import save_image
+from torch.utils.data import DataLoader
+from dataset import MNISTDataModule
+import pytorch_lightning as pl
+from model import VAEpl
+from pytorch_lightning.loggers import TensorBoardLogger
+from pytorch_lightning.strategies import DeepSpeedStrategy
+torch.set_float32_matmul_precision("medium")
+
+""" 
+GOALS:
+* Understand the strategy (deepspeed, ddp, etc) and how to use it
+* Setup a config, for scheduler etc instead of configuring it in each sub-module
+* Metrics
+"""
+
+
+# things to add
+lr = 3e-4
+batch_size = 128
+num_workers = 2
+model = VAEpl(lr)
+dm = MNISTDataModule(batch_size, num_workers)
+logger = TensorBoardLogger("my_checkpoint", name="scheduler_autolr_vae_pl_model")
+
+# add callback for learning rate monitor, model checkpoint, and scheduler on plateau
+callbacks = [pl.callbacks.LearningRateMonitor(logging_interval="step"),
+             pl.callbacks.ModelCheckpoint(monitor="val_loss", save_top_k=1, mode="min", save_last=True),
+             ]
+
+if __name__ == "__main__":
+    trainer = pl.Trainer(
+        max_epochs=100,
+        accelerator="gpu",
+        devices=2,
+        logger=logger,
+        #precision=16,
+        strategy=DeepSpeedStrategy(
+            stage=0,
+        ),
+    )
+
+    #trainer.tune(model, dm)
+    trainer.fit(model, dm) 

ML/Pytorch/more_advanced/VAE/lightning_vae/utils.py

@@ -0,0 +1,41 @@
+import torch 
+import torch.nn as nn
+import torch.nn.functional as F
+# import save_image from torchvision.utils
+from torchvision.utils import save_image
+
+
+def inference(model, dataset, digit, num_examples=1):
+    """
+    Generates (num_examples) of a particular digit.
+    Specifically we extract an example of each digit,
+    then after we have the mu, sigma representation for
+    each digit we can sample from that.
+
+    After we sample we can run the decoder part of the VAE
+    and generate examples.
+    """
+    images = []
+    idx = 0
+    for x, y in dataset:
+        if y == idx:
+            images.append(x)
+            idx += 1
+        if idx == 10:
+            break
+
+    encodings_digit = []
+    for d in range(10):
+        with torch.no_grad():
+            mu, sigma = model.encode(images[d].view(1, 784))
+        encodings_digit.append((mu, sigma))
+
+    mu, sigma = encodings_digit[digit]
+    for example in range(num_examples):
+        epsilon = torch.randn_like(sigma)
+        z = mu + sigma * epsilon
+        out = model.decode(z)
+        out = out.view(-1, 1, 28, 28)
+        save_image(out, f"generated_{digit}_ex{example}.png")
+
+

ML/Pytorch/more_advanced/VAE/train.py

@@ -23,27 +23,7 @@ model = VariationalAutoEncoder(INPUT_DIM, H_DIM, Z_DIM).to(DEVICE)
 optimizer = optim.Adam(model.parameters(), lr=LR_RATE)
 loss_fn = nn.BCELoss(reduction="sum")
 
-# Start Training
-for epoch in range(NUM_EPOCHS):
-    loop = tqdm(enumerate(train_loader))
-    for i, (x, _) in loop:
-        # Forward pass
-        x = x.to(DEVICE).view(x.shape[0], INPUT_DIM)
-        x_reconstructed, mu, sigma = model(x)
 
-        # Compute loss
-        reconstruction_loss = loss_fn(x_reconstructed, x)
-        kl_div = -torch.sum(1 + torch.log(sigma.pow(2)) - mu.pow(2) - sigma.pow(2))
-
-        # Backprop
-        loss = reconstruction_loss + kl_div
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-        loop.set_postfix(loss=loss.item())
-
-
-model = model.to("cpu")
 def inference(digit, num_examples=1):
     """
     Generates (num_examples) of a particular digit.
@@ -79,8 +59,3 @@ def inference(digit, num_examples=1):
 
 for idx in range(10):
     inference(idx, num_examples=5)
-
-
-
-
-