normalizing flows incomplete

course_UvA-DL/09-normalizing-flows/NF_image_modeling.py

@@ -13,7 +13,7 @@
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
-import pytorch_lightning as pl
+import lightning as L
 import seaborn as sns
 import tabulate
 import torch
@@ -22,16 +22,17 @@
 import torch.optim as optim
 import torch.utils.data as data
 import torchvision
-from IPython.display import HTML, display, set_matplotlib_formats
+from IPython.display import HTML, display
+import matplotlib_inline.backend_inline
 from matplotlib.colors import to_rgb
-from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
+from lightning.pytorch.callbacks import LearningRateMonitor, ModelCheckpoint
 from torch import Tensor
 from torchvision import transforms
 from torchvision.datasets import MNIST
 from tqdm.notebook import tqdm
 
 # %matplotlib inline
-set_matplotlib_formats("svg", "pdf")  # For export
+matplotlib_inline.backend_inline.set_matplotlib_formats("svg", "pdf")  # For export
 matplotlib.rcParams["lines.linewidth"] = 2.0
 sns.reset_orig()
 
@@ -41,7 +42,7 @@
 CHECKPOINT_PATH = os.environ.get("PATH_CHECKPOINT", "saved_models/tutorial11")
 
 # Setting the seed
-pl.seed_everything(42)
+L.seed_everything(42)
 
 # Ensure that all operations are deterministic on GPU (if used) for reproducibility
 torch.backends.cudnn.determinstic = True
@@ -97,7 +98,7 @@ def discretize(sample):
 
 # Loading the training dataset. We need to split it into a training and validation part
 train_dataset = MNIST(root=DATASET_PATH, train=True, transform=transform, download=True)
-pl.seed_everything(42)
+L.seed_everything(42)
 train_set, val_set = torch.utils.data.random_split(train_dataset, [50000, 10000])
 
 # Loading the test set
@@ -258,7 +259,7 @@ def show_imgs(imgs, title=None, row_size=4):
 
 
 # %%
-class ImageFlow(pl.LightningModule):
+class ImageFlow(L.LightningModule):
     def __init__(self, flows, import_samples=8):
         """
         Args:
@@ -449,7 +450,7 @@ def dequant(self, z, ldj):
 
 # %%
 # Testing invertibility of dequantization layer
-pl.seed_everything(42)
+L.seed_everything(42)
 orig_img = train_set[0][0].unsqueeze(dim=0)
 ldj = torch.zeros(
     1,
@@ -916,9 +917,10 @@ def create_simple_flow(use_vardeq=True):
 # %%
 def train_flow(flow, model_name="MNISTFlow"):
     # Create a PyTorch Lightning trainer
-    trainer = pl.Trainer(
+    trainer = L.Trainer(
         default_root_dir=os.path.join(CHECKPOINT_PATH, model_name),
-        gpus=1 if torch.cuda.is_available() else 0,
+        accelerator="auto",
+        devices=1,
         max_epochs=200,
         gradient_clip_val=1.0,
         callbacks=[
@@ -1216,12 +1218,12 @@ def print_num_params(model):
 # The seeds are set to obtain reproducable generations and are not cherry picked.
 
 # %%
-pl.seed_everything(44)
+L.seed_everything(44)
 samples = flow_dict["vardeq"]["model"].sample(img_shape=[16, 1, 28, 28])
 show_imgs(samples.cpu())
 
 # %%
-pl.seed_everything(44)
+L.seed_everything(44)
 samples = flow_dict["multiscale"]["model"].sample(img_shape=[16, 8, 7, 7])
 show_imgs(samples.cpu())
 
@@ -1262,12 +1264,12 @@ def interpolate(model, img1, img2, num_steps=8):
 exmp_imgs, _ = next(iter(train_loader))
 
 # %%
-pl.seed_everything(42)
+L.seed_everything(42)
 for i in range(2):
     interpolate(flow_dict["vardeq"]["model"], exmp_imgs[2 * i], exmp_imgs[2 * i + 1])
 
 # %%
-pl.seed_everything(42)
+L.seed_everything(42)
 for i in range(2):
     interpolate(flow_dict["multiscale"]["model"], exmp_imgs[2 * i], exmp_imgs[2 * i + 1])
 
@@ -1290,7 +1292,7 @@ def interpolate(model, img1, img2, num_steps=8):
 # Below we visualize three examples of this:
 
 # %%
-pl.seed_everything(44)
+L.seed_everything(44)
 for _ in range(3):
     z_init = flow_dict["multiscale"]["model"].prior.sample(sample_shape=[1, 8, 7, 7])
     z_init = z_init.expand(8, -1, -1, -1)