Tutorial 3: Image, Conditional Diffusion and Beyond
Contents
Tutorial 3: Image, Conditional Diffusion and Beyond¶
Week 2, Day 4: Name of the day
By Neuromatch Academy
Content creators: Binxu Wang
Content reviewers: Shaonan Wang, Dongrui Deng, Dora Zhiyu Yang, Adrita Das
Content editors: Shaonan Wang
Production editors: Spiros Chavlis
Tutorial Objectives¶
Understand the idea behind Diffusion generative models: score and reversal of diffusion process.
Learn the score function by denoising data.
Hands-on experience in learning the score to generate certain distributions.
Setup¶
⚠ Experimental LLM-enhanced tutorial ⚠
This notebook includes Neuromatch’s experimental Chatify 🤖 functionality. The Chatify notebook extension adds support for a large language model-based “coding tutor” to the materials. The tutor provides automatically generated text to help explain any code cell in this notebook.
Note that using Chatify may cause breaking changes and/or provide incorrect or misleading information. If you wish to proceed by installing and enabling the Chatify extension, you should run the next two code blocks (hidden by default). If you do not want to use this experimental version of the Neuromatch materials, please use the stable materials instead.
To use the Chatify helper, insert the %%explain magic command at the start of any code cell and then run it (shift + enter) to access an interface for receiving LLM-based assitance. You can then select different options from the dropdown menus depending on what sort of assitance you want. To disable Chatify and run the code block as usual, simply delete the %%explain command and re-run the cell.
Note that, by default, all of Chatify’s responses are generated locally. This often takes several minutes per response. Once you click the “Submit request” button, just be patient– stuff is happening even if you can’t see it right away!
Thanks for giving Chatify a try! Love it? Hate it? Either way, we’d love to hear from you about your Chatify experience! Please consider filling out our brief survey to provide feedback and help us make Chatify more awesome!
Run the next two cells to install and configure Chatify…
%pip install -q davos
import davos
davos.config.suppress_stdout = True
Note: you may need to restart the kernel to use updated packages.
smuggle chatify # pip: git+https://github.com/ContextLab/chatify.git
%load_ext chatify
Using default configuration!
Downloading the 'cache' file.
Install dependencies¶
WARNING: There may be errors and/or warnings reported during the installation. However, they are to be ignored.
# @title Install dependencies
# @markdown **WARNING**: There may be *errors* and/or *warnings* reported during the installation. However, they are to be ignored.
!pip install pillow --quiet
!pip install diffusers transformers tokenizers --quiet
!pip install accelerate --quiet
Install and import feedback gadget¶
# @title Install and import feedback gadget
!pip3 install vibecheck datatops --quiet
from vibecheck import DatatopsContentReviewContainer
def content_review(notebook_section: str):
return DatatopsContentReviewContainer(
"", # No text prompt
notebook_section,
{
"url": "https://pmyvdlilci.execute-api.us-east-1.amazonaws.com/klab",
"name": "neuromatch_dl",
"user_key": "f379rz8y",
},
).render()
feedback_prefix = "W2D4_T3"
# Imports
import random
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import functools
from torch.optim import Adam
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision.datasets import MNIST
from tqdm.notebook import trange, tqdm
from torch.optim.lr_scheduler import MultiplicativeLR, LambdaLR
from torchvision.utils import make_grid
Figure settings¶
# @title Figure settings
import ipywidgets as widgets # interactive display
%config InlineBackend.figure_format = 'retina'
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/content-creation/main/nma.mplstyle")
Set random seed¶
Executing set_seed(seed=seed) you are setting the seed
# @title Set random seed
# @markdown Executing `set_seed(seed=seed)` you are setting the seed
# For DL its critical to set the random seed so that students can have a
# baseline to compare their results to expected results.
# Read more here: https://pytorch.org/docs/stable/notes/randomness.html
# Call `set_seed` function in the exercises to ensure reproducibility.
import random
import torch
def set_seed(seed=None, seed_torch=True):
"""
Function that controls randomness.
NumPy and random modules must be imported.
Args:
seed : Integer
A non-negative integer that defines the random state. Default is `None`.
seed_torch : Boolean
If `True` sets the random seed for pytorch tensors, so pytorch module
must be imported. Default is `True`.
Returns:
Nothing.
"""
if seed is None:
seed = np.random.choice(2 ** 32)
random.seed(seed)
np.random.seed(seed)
if seed_torch:
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.cuda.manual_seed(seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
print(f'Random seed {seed} has been set.')
# In case that `DataLoader` is used
def seed_worker(worker_id):
"""
DataLoader will reseed workers following randomness in
multi-process data loading algorithm.
Args:
worker_id: integer
ID of subprocess to seed. 0 means that
the data will be loaded in the main process
Refer: https://pytorch.org/docs/stable/data.html#data-loading-randomness for more details
Returns:
Nothing
"""
worker_seed = torch.initial_seed() % 2**32
np.random.seed(worker_seed)
random.seed(worker_seed)
Set device (GPU or CPU). Execute set_device()¶
# @title Set device (GPU or CPU). Execute `set_device()`
# Inform the user if the notebook uses GPU or CPU.
def set_device():
"""
Set the device. CUDA if available, CPU otherwise
Args:
None
Returns:
Nothing
"""
device = "cuda" if torch.cuda.is_available() else "cpu"
if device != "cuda":
print("WARNING: For this notebook to perform best, "
"if possible, in the menu under `Runtime` -> "
"`Change runtime type.` select `GPU` ")
else:
print("GPU is enabled in this notebook.")
return device
DEVICE = set_device()
SEED = 2021
set_seed(seed=SEED)
WARNING: For this notebook to perform best, if possible, in the menu under `Runtime` -> `Change runtime type.` select `GPU`
Random seed 2021 has been set.
Neural Network Architecture¶
We just learned the basic principles of diffusion models, with the takeaway that the score function allows us to turn pure noise into some interesting data distribution. Further, we will approximate the score function with a neural network via denoising score matching. But when working with images, we need our neural network to ‘play nice’ with them and to reflect the inductive biases we associate with images.
A reasonable choice is to choose the neural network architecture to be that of a U-Net, which is a CNN-like architecture with:
downscaling/upscaling operations that help the network process features of images at different spatial scales.
skip connection as an information highway.
Since the score function we’re trying to learn a function of time, we also need to devise a way to ensure our neural network properly responds to changes in time. For this purpose, we can use a time embedding.
Video 1: Network architecture¶
Submit your feedback¶
# @title Submit your feedback
content_review(f"{feedback_prefix}_Network_Architecture_Video")
Coding Exercise 1: Train Diffusion for MNIST¶
Finally, let’s implement and train an actual image diffusion model for the MNIST dataset.
By examining the neural network architecture of the score approximator, you will understand the inductive biases we built in.
In the next cell, you will implement the helper functions for the forward process.
marginal_prob_stdfor \(\sigma_t\) (note, it’s standard deviation, not the variance)diffusion_coefffor \(g(t)\)
Math Recap for Forward Processes:
We will use the same forward process (variance exploding SDE) as in the last tutorial, which reads:
and we let the diffusion coefficient \(g(t)=\lambda^t\), with \(\lambda > 1\).
If so, the marginal distribution of state \(\mathbf x_t\) at time t given an initial state \(\mathbf x_0\) will be a Gaussian \(\mathcal N(\mathbf x_t|\mathbf x_0,\sigma_t^2 I)\). The variance is the integration of the squared diffusion coefficient.
def marginal_prob_std(t, Lambda, device='cpu'):
"""Compute the standard deviation of $p_{0t}(x(t) | x(0))$.
Args:
t: A vector of time steps.
Lambda: The $\lambda$ in our SDE.
Returns:
std : The standard deviation.
"""
t = t.to(device)
#################################################
## TODO for students: Implement the standard deviation
raise NotImplementedError("Student exercise: Implement the standard deviation")
#################################################
std = ...
return std
def diffusion_coeff(t, Lambda, device='cpu'):
"""Compute the diffusion coefficient of our SDE.
Args:
t: A vector of time steps.
Lambda: The $\lambda$ in our SDE.
Returns:
diff_coeff : The vector of diffusion coefficients.
"""
#################################################
## TODO for students: Implement the diffusion coefficients
raise NotImplementedError("Student exercise: Implement the diffusion coefficients")
#################################################
diff_coeff = ...
return diff_coeff.to(device)
Submit your feedback¶
# @title Submit your feedback
content_review(f"{feedback_prefix}_Train_Diffusion_for_MNIST_Exercise")
Network architecture¶
Below is code for a simple time embedding and modulation layer. Basically, time \(t\) is multiplexed as sine and cosine basis, then a linear readout creates the time modulation signal.
Time embedding and modulation¶
# @title Time embedding and modulation
class GaussianFourierProjection(nn.Module):
"""Gaussian random features for encoding time steps."""
def __init__(self, embed_dim, scale=30.):
super().__init__()
# Randomly sample weights (frequencies) during initialization.
# These weights (frequencies) are fixed during optimization and are not trainable.
self.W = nn.Parameter(torch.randn(embed_dim // 2) * scale, requires_grad=False)
def forward(self, x):
# Cosine(2 pi freq x), Sine(2 pi freq x)
x_proj = x[:, None] * self.W[None, :] * 2 * np.pi
return torch.cat([torch.sin(x_proj), torch.cos(x_proj)], dim=-1)
class Dense(nn.Module):
"""A fully connected layer that reshapes outputs to feature maps.
Allow time repr to input additively from the side of a convolution layer.
"""
def __init__(self, input_dim, output_dim):
super().__init__()
self.dense = nn.Linear(input_dim, output_dim)
def forward(self, x):
# this broadcast the 2d tensor to 4d, add the same value across space.
return self.dense(x)[..., None, None]
Below is code for a simple U-Net architecture. Apparently, diffusion models can be more or less successful with different architectural details. So this example is mainly for illustrative purposes.
Time-dependent UNet score model¶
# @title Time-dependent UNet score model
class UNet(nn.Module):
"""A time-dependent score-based model built upon U-Net architecture."""
def __init__(self, marginal_prob_std, channels=[32, 64, 128, 256], embed_dim=256):
"""Initialize a time-dependent score-based network.
Args:
marginal_prob_std: A function that takes time t and gives the standard
deviation of the perturbation kernel p_{0t}(x(t) | x(0)).
channels: The number of channels for feature maps of each resolution.
embed_dim: The dimensionality of Gaussian random feature embeddings.
"""
super().__init__()
# Gaussian random feature embedding layer for time
self.time_embed = nn.Sequential(
GaussianFourierProjection(embed_dim=embed_dim),
nn.Linear(embed_dim, embed_dim)
)
# Encoding layers where the resolution decreases
self.conv1 = nn.Conv2d(1, channels[0], 3, stride=1, bias=False)
self.t_mod1 = Dense(embed_dim, channels[0])
self.gnorm1 = nn.GroupNorm(4, num_channels=channels[0])
self.conv2 = nn.Conv2d(channels[0], channels[1], 3, stride=2, bias=False)
self.t_mod2 = Dense(embed_dim, channels[1])
self.gnorm2 = nn.GroupNorm(32, num_channels=channels[1])
self.conv3 = nn.Conv2d(channels[1], channels[2], 3, stride=2, bias=False)
self.t_mod3 = Dense(embed_dim, channels[2])
self.gnorm3 = nn.GroupNorm(32, num_channels=channels[2])
self.conv4 = nn.Conv2d(channels[2], channels[3], 3, stride=2, bias=False)
self.t_mod4 = Dense(embed_dim, channels[3])
self.gnorm4 = nn.GroupNorm(32, num_channels=channels[3])
# Decoding layers where the resolution increases
self.tconv4 = nn.ConvTranspose2d(channels[3], channels[2], 3, stride=2, bias=False)
self.t_mod5 = Dense(embed_dim, channels[2])
self.tgnorm4 = nn.GroupNorm(32, num_channels=channels[2])
self.tconv3 = nn.ConvTranspose2d(channels[2] + channels[2], channels[1], 3, stride=2, bias=False, output_padding=1)
self.t_mod6 = Dense(embed_dim, channels[1])
self.tgnorm3 = nn.GroupNorm(32, num_channels=channels[1])
self.tconv2 = nn.ConvTranspose2d(channels[1] + channels[1], channels[0], 3, stride=2, bias=False, output_padding=1)
self.t_mod7 = Dense(embed_dim, channels[0])
self.tgnorm2 = nn.GroupNorm(32, num_channels=channels[0])
self.tconv1 = nn.ConvTranspose2d(channels[0] + channels[0], 1, 3, stride=1)
# The swish activation function
self.act = lambda x: x * torch.sigmoid(x)
# A restricted version of the `marginal_prob_std` function, after specifying a Lambda.
self.marginal_prob_std = marginal_prob_std
def forward(self, x, t, y=None):
# Obtain the Gaussian random feature embedding for t
embed = self.act(self.time_embed(t))
# Encoding path, downsampling
## Incorporate information from t
h1 = self.conv1(x) + self.t_mod1(embed)
## Group normalization and apply activation function
h1 = self.act(self.gnorm1(h1))
# 2nd conv
h2 = self.conv2(h1) + self.t_mod2(embed)
h2 = self.act(self.gnorm2(h2))
# 3rd conv
h3 = self.conv3(h2) + self.t_mod3(embed)
h3 = self.act(self.gnorm3(h3))
# 4th conv
h4 = self.conv4(h3) + self.t_mod4(embed)
h4 = self.act(self.gnorm4(h4))
# Decoding path up sampling
h = self.tconv4(h4) + self.t_mod5(embed)
## Skip connection from the encoding path
h = self.act(self.tgnorm4(h))
h = self.tconv3(torch.cat([h, h3], dim=1)) + self.t_mod6(embed)
h = self.act(self.tgnorm3(h))
h = self.tconv2(torch.cat([h, h2], dim=1)) + self.t_mod7(embed)
h = self.act(self.tgnorm2(h))
h = self.tconv1(torch.cat([h, h1], dim=1))
# Normalize output
h = h / self.marginal_prob_std(t)[:, None, None, None]
return h
Think! 1: U-Net Architecture¶
Looking at the U-Net architecture, can you find the module(s) corresponding to the following operations?
Downsampling the spatial features?
Upsampling the spatial features?
The skip connection from the down branch to the up branch, how is it implemented?
How is time modulation implemented?
Why is the output divided by
self.marginal_prob_std(t)before output? How might this help or harm the score learning?
Take 2 minutes to think in silence, then discuss as a group (~10 minutes).
Submit your feedback¶
# @title Submit your feedback
content_review(f"{feedback_prefix}_UNet_Architecture_Discussion")
Coding Exercise 2: Defining the loss function¶
In the next cell, you will implement the denoising score matching (DSM) objective as we used in the last tutorial.
where the time weighting is chosen as \(\gamma_t=\sigma_t^2\), which emphasizes the high noise period (\(t\sim 1\)) more than the low noise period (\(t\sim 0\)).
Tips:
The major difference from the last tutorial is that the score \(s\), noise \(z\), and states \(x\) are all batch image-shaped tensor, so remember to broadcast the \(\sigma_t\) properly. e.g. this
std[:, None, None, None]will be helpful.epsis set at a small number to stop the model from learning the score function of a very small noise scale, which is highly irregular.
def loss_fn(model, x, marginal_prob_std, eps=1e-3, device='cpu'):
"""The loss function for training score-based generative models.
Args:
model: A PyTorch model instance that represents a
time-dependent score-based model.
Note, it takes two inputs in its forward function model(x, t)
$s_\theta(x,t)$ in the equation
x: A mini-batch of training data.
marginal_prob_std: A function that gives the standard deviation of
the perturbation kernel, takes `t` as input.
$\sigma_t$ in the equation.
eps: A tolerance value for numerical stability.
"""
# Sample time uniformly in eps, 1
random_t = torch.rand(x.shape[0], device=device) * (1. - eps) + eps
# Find the noise std at the time `t`
std = marginal_prob_std(random_t).to(device)
#################################################
## TODO for students: Implement the denoising score matching eq.
raise NotImplementedError("Student exercise: Implement the denoising score matching eq. ")
#################################################
# get normally distributed noise N(0, I)
z = ...
# compute the perturbed x = x + z * \sigma_t
perturbed_x = ...
# predict score with the model at (perturbed x, t)
score = ...
# compute distance between the score and noise \| score * sigma_t + z \|_2^2
loss = ...
##############
return loss
A correctly implemented loss function shall pass the test below.
For a dataset with a single 0 datapoint, we have the analytical score \(\mathbf s(\mathbf x,t)=-\mathbf x/\sigma_t^2\). We test that, for this case, the analytical has zero loss.
Test loss function¶
# @title Test loss function
marginal_prob_std_test = lambda t: marginal_prob_std(t, Lambda=10, device='cpu')
score_analyt_test = lambda x_t, t: - x_t / marginal_prob_std_test(t)[:,None,None,None]**2
x_test = torch.zeros(10, 3, 64, 64)
loss = loss_fn(score_analyt_test, x_test, marginal_prob_std_test, eps=1e-3, device='cpu')
assert torch.allclose(loss,torch.zeros(1)), "the loss should be zero in this case"
Submit your feedback¶
# @title Submit your feedback
content_review(f"{feedback_prefix}_Defining_the_loss_function_Exercise")
Train and Test the Diffusion Model¶
Note: We have reduced the n_epochs to 12, but feel free to increase and use a larger value. The original value was set to 100, but if the training takes too long, n_epochs=50 with batch_size=1024 also suffice. An average loss of around ~30 can generate acceptable digits.
Training the model¶
# @title Training the model
Lambda = 25.0 # @param {'type':'number'}
marginal_prob_std_fn = lambda t: marginal_prob_std(t, Lambda=Lambda, device=DEVICE)
diffusion_coeff_fn = lambda t: diffusion_coeff(t, Lambda=Lambda, device=DEVICE)
score_model = UNet(marginal_prob_std=marginal_prob_std_fn)
score_model = score_model.to(DEVICE)
n_epochs = 12 # @param {'type':'integer'}
# size of a mini-batch
batch_size = 1024 # @param {'type':'integer'}
# learning rate
lr = 10e-4 # @param {'type':'number'}
set_seed(SEED)
dataset = MNIST('.', train=True, transform=transforms.ToTensor(), download=True)
g = torch.Generator()
g.manual_seed(SEED)
data_loader = DataLoader(dataset, batch_size=batch_size,
shuffle=True, num_workers=2,
worker_init_fn=seed_worker,
generator=g,)
optimizer = Adam(score_model.parameters(), lr=lr)
scheduler = LambdaLR(optimizer, lr_lambda=lambda epoch: max(0.2, 1 - epoch / n_epochs))
tqdm_epoch = trange(n_epochs)
for epoch in tqdm_epoch:
avg_loss = 0.
num_items = 0
pbar = tqdm(data_loader)
for x, y in pbar:
x = x.to(DEVICE)
loss = loss_fn(score_model, x, marginal_prob_std_fn, eps=0.01, device=DEVICE)
optimizer.zero_grad()
loss.backward()
optimizer.step()
avg_loss += loss.item() * x.shape[0]
num_items += x.shape[0]
scheduler.step()
print(f"Average Loss: {(avg_loss / num_items):5f} lr {scheduler.get_last_lr()[0]:.1e}")
# Print the averaged training loss so far.
tqdm_epoch.set_description(f'Average Loss: {(avg_loss / num_items):.5f}')
# Update the checkpoint after each epoch of training.
torch.save(score_model.state_dict(), 'ckpt.pth')
Random seed 2021 has been set.
Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to ./MNIST/raw/train-images-idx3-ubyte.gz
Extracting ./MNIST/raw/train-images-idx3-ubyte.gz to ./MNIST/raw
Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ./MNIST/raw/train-labels-idx1-ubyte.gz
Extracting ./MNIST/raw/train-labels-idx1-ubyte.gz to ./MNIST/raw
Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ./MNIST/raw/t10k-images-idx3-ubyte.gz
Extracting ./MNIST/raw/t10k-images-idx3-ubyte.gz to ./MNIST/raw
Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ./MNIST/raw/t10k-labels-idx1-ubyte.gz
Extracting ./MNIST/raw/t10k-labels-idx1-ubyte.gz to ./MNIST/raw
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
File ~/opt/anaconda3/envs/nma-course/lib/python3.9/site-packages/torch/utils/data/dataloader.py:1132, in _MultiProcessingDataLoaderIter._try_get_data(self, timeout)
1131 try:
-> 1132 data = self._data_queue.get(timeout=timeout)
1133 return (True, data)
File ~/opt/anaconda3/envs/nma-course/lib/python3.9/multiprocessing/queues.py:113, in Queue.get(self, block, timeout)
112 timeout = deadline - time.monotonic()
--> 113 if not self._poll(timeout):
114 raise Empty
File ~/opt/anaconda3/envs/nma-course/lib/python3.9/multiprocessing/connection.py:257, in _ConnectionBase.poll(self, timeout)
256 self._check_readable()
--> 257 return self._poll(timeout)
File ~/opt/anaconda3/envs/nma-course/lib/python3.9/multiprocessing/connection.py:424, in Connection._poll(self, timeout)
423 def _poll(self, timeout):
--> 424 r = wait([self], timeout)
425 return bool(r)
File ~/opt/anaconda3/envs/nma-course/lib/python3.9/multiprocessing/connection.py:931, in wait(object_list, timeout)
930 while True:
--> 931 ready = selector.select(timeout)
932 if ready:
File ~/opt/anaconda3/envs/nma-course/lib/python3.9/selectors.py:416, in _PollLikeSelector.select(self, timeout)
415 try:
--> 416 fd_event_list = self._selector.poll(timeout)
417 except InterruptedError:
File ~/opt/anaconda3/envs/nma-course/lib/python3.9/site-packages/torch/utils/data/_utils/signal_handling.py:66, in _set_SIGCHLD_handler.<locals>.handler(signum, frame)
63 def handler(signum, frame):
64 # This following call uses `waitid` with WNOHANG from C side. Therefore,
65 # Python can still get and update the process status successfully.
---> 66 _error_if_any_worker_fails()
67 if previous_handler is not None:
RuntimeError: DataLoader worker (pid 80683) exited unexpectedly with exit code 1. Details are lost due to multiprocessing. Rerunning with num_workers=0 may give better error trace.
The above exception was the direct cause of the following exception:
RuntimeError Traceback (most recent call last)
Cell In[21], line 32
30 num_items = 0
31 pbar = tqdm(data_loader)
---> 32 for x, y in pbar:
33 x = x.to(DEVICE)
34 loss = loss_fn(score_model, x, marginal_prob_std_fn, eps=0.01, device=DEVICE)
File ~/opt/anaconda3/envs/nma-course/lib/python3.9/site-packages/tqdm/notebook.py:254, in tqdm_notebook.__iter__(self)
252 try:
253 it = super(tqdm_notebook, self).__iter__()
--> 254 for obj in it:
255 # return super(tqdm...) will not catch exception
256 yield obj
257 # NB: except ... [ as ...] breaks IPython async KeyboardInterrupt
File ~/opt/anaconda3/envs/nma-course/lib/python3.9/site-packages/tqdm/std.py:1178, in tqdm.__iter__(self)
1175 time = self._time
1177 try:
-> 1178 for obj in iterable:
1179 yield obj
1180 # Update and possibly print the progressbar.
1181 # Note: does not call self.update(1) for speed optimisation.
File ~/opt/anaconda3/envs/nma-course/lib/python3.9/site-packages/torch/utils/data/dataloader.py:633, in _BaseDataLoaderIter.__next__(self)
630 if self._sampler_iter is None:
631 # TODO(https://github.com/pytorch/pytorch/issues/76750)
632 self._reset() # type: ignore[call-arg]
--> 633 data = self._next_data()
634 self._num_yielded += 1
635 if self._dataset_kind == _DatasetKind.Iterable and \
636 self._IterableDataset_len_called is not None and \
637 self._num_yielded > self._IterableDataset_len_called:
File ~/opt/anaconda3/envs/nma-course/lib/python3.9/site-packages/torch/utils/data/dataloader.py:1328, in _MultiProcessingDataLoaderIter._next_data(self)
1325 return self._process_data(data)
1327 assert not self._shutdown and self._tasks_outstanding > 0
-> 1328 idx, data = self._get_data()
1329 self._tasks_outstanding -= 1
1330 if self._dataset_kind == _DatasetKind.Iterable:
1331 # Check for _IterableDatasetStopIteration
File ~/opt/anaconda3/envs/nma-course/lib/python3.9/site-packages/torch/utils/data/dataloader.py:1294, in _MultiProcessingDataLoaderIter._get_data(self)
1290 # In this case, `self._data_queue` is a `queue.Queue`,. But we don't
1291 # need to call `.task_done()` because we don't use `.join()`.
1292 else:
1293 while True:
-> 1294 success, data = self._try_get_data()
1295 if success:
1296 return data
File ~/opt/anaconda3/envs/nma-course/lib/python3.9/site-packages/torch/utils/data/dataloader.py:1145, in _MultiProcessingDataLoaderIter._try_get_data(self, timeout)
1143 if len(failed_workers) > 0:
1144 pids_str = ', '.join(str(w.pid) for w in failed_workers)
-> 1145 raise RuntimeError('DataLoader worker (pid(s) {}) exited unexpectedly'.format(pids_str)) from e
1146 if isinstance(e, queue.Empty):
1147 return (False, None)
RuntimeError: DataLoader worker (pid(s) 80683) exited unexpectedly
Define the Sampler¶
# @title Define the Sampler
def Euler_Maruyama_sampler(score_model,
marginal_prob_std,
diffusion_coeff,
batch_size=64,
x_shape=(1, 28, 28),
num_steps=500,
device='cuda',
eps=1e-3, y=None):
"""Generate samples from score-based models with the Euler-Maruyama solver.
Args:
score_model: A PyTorch model that represents the time-dependent score-based model.
marginal_prob_std: A function that gives the standard deviation of
the perturbation kernel.
diffusion_coeff: A function that gives the diffusion coefficient of the SDE.
batch_size: The number of samplers to generate by calling this function once.
num_steps: The number of sampling steps.
Equivalent to the number of discretized time steps.
device: 'cuda' for running on GPUs, and 'cpu' for running on CPUs.
eps: The smallest time step for numerical stability.
Returns:
Samples.
"""
t = torch.ones(batch_size).to(device)
r = torch.randn(batch_size, *x_shape).to(device)
init_x = r * marginal_prob_std(t)[:, None, None, None]
init_x = init_x.to(device)
time_steps = torch.linspace(1., eps, num_steps).to(device)
step_size = time_steps[0] - time_steps[1]
x = init_x
with torch.no_grad():
for time_step in tqdm(time_steps):
batch_time_step = torch.ones(batch_size, device=device) * time_step
g = diffusion_coeff(batch_time_step)
mean_x = x + (g**2)[:, None, None, None] * score_model(x, batch_time_step, y=y) * step_size
x = mean_x + torch.sqrt(step_size) * g[:, None, None, None] * torch.randn_like(x)
# Do not include any noise in the last sampling step.
return mean_x
Sampling¶
# @title Sampling
def save_samples_uncond(score_model, suffix="", device='cpu'):
score_model.eval()
## Generate samples using the specified sampler.
sample_batch_size = 64 # @param {'type':'integer'}
num_steps = 250 # @param {'type':'integer'}
# score_model.eval()
## Generate samples using the specified sampler.
samples = Euler_Maruyama_sampler(score_model,
marginal_prob_std_fn,
diffusion_coeff_fn,
sample_batch_size,
num_steps=num_steps,
device=DEVICE,
eps=0.001)
# Sample visualization.
samples = samples.clamp(0.0, 1.0)
sample_grid = make_grid(samples, nrow=int(np.sqrt(sample_batch_size)))
sample_np = sample_grid.permute(1, 2, 0).cpu().numpy()
plt.imsave(f"uncondition_diffusion{suffix}.png", sample_np, )
plt.figure(figsize=(6,6))
plt.axis('off')
plt.imshow(sample_np, vmin=0., vmax=1.)
plt.show()
marginal_prob_std_fn = lambda t: marginal_prob_std(t, Lambda=Lambda, device=DEVICE)
uncond_score_model = UNet(marginal_prob_std=marginal_prob_std_fn)
uncond_score_model.load_state_dict(torch.load("ckpt.pth"))
uncond_score_model.to(DEVICE)
save_samples_uncond(uncond_score_model, suffix="", device=DEVICE)
---------------------------------------------------------------------------
FileNotFoundError Traceback (most recent call last)
Cell In[23], line 30
28 marginal_prob_std_fn = lambda t: marginal_prob_std(t, Lambda=Lambda, device=DEVICE)
29 uncond_score_model = UNet(marginal_prob_std=marginal_prob_std_fn)
---> 30 uncond_score_model.load_state_dict(torch.load("ckpt.pth"))
31 uncond_score_model.to(DEVICE)
32 save_samples_uncond(uncond_score_model, suffix="", device=DEVICE)
File ~/opt/anaconda3/envs/nma-course/lib/python3.9/site-packages/torch/serialization.py:791, in load(f, map_location, pickle_module, weights_only, **pickle_load_args)
788 if 'encoding' not in pickle_load_args.keys():
789 pickle_load_args['encoding'] = 'utf-8'
--> 791 with _open_file_like(f, 'rb') as opened_file:
792 if _is_zipfile(opened_file):
793 # The zipfile reader is going to advance the current file position.
794 # If we want to actually tail call to torch.jit.load, we need to
795 # reset back to the original position.
796 orig_position = opened_file.tell()
File ~/opt/anaconda3/envs/nma-course/lib/python3.9/site-packages/torch/serialization.py:271, in _open_file_like(name_or_buffer, mode)
269 def _open_file_like(name_or_buffer, mode):
270 if _is_path(name_or_buffer):
--> 271 return _open_file(name_or_buffer, mode)
272 else:
273 if 'w' in mode:
File ~/opt/anaconda3/envs/nma-course/lib/python3.9/site-packages/torch/serialization.py:252, in _open_file.__init__(self, name, mode)
251 def __init__(self, name, mode):
--> 252 super().__init__(open(name, mode))
FileNotFoundError: [Errno 2] No such file or directory: 'ckpt.pth'
Nice job! You have just finished the training of a Diffusion model. As you see, the result is not ideal, and many factors affect this. To name a few:
Better network architecture: residual connections, attention mechanism, better upsampling mechanism
Better objective: better weighting function \(\gamma_t\)
Better optimization procedure: using learning rate decay
Better sampling algorithm: Euler integration is known to have larger errors, so it’s advisable to use a more advanced SDE or ODE solver
Section 2: Conditional Diffusion Model¶
Another way to greatly improve the result is adding a conditional signal – for example, tell the score network which digit you want. This makes the score modeling much more effortless and adds controllability to the user. The popular Stable Diffusion model is one of this kind, which uses natural language text as the conditional signal for images.
Video 2: Conditional Diffusion Model¶
Submit your feedback¶
# @title Submit your feedback
content_review(f"{feedback_prefix}_Conditional_Diffusion_Model_Video")
In formulation, the conditional diffusion is highly similar to the unconditional diffusion.
If you are curious about how to build and train a conditional diffusion model, you are welcome to look at the Bonus exercise at the end.
Video 3: Advanced Techinque - Stable Diffusion¶
Submit your feedback¶
# @title Submit your feedback
content_review(f"{feedback_prefix}_Advanced_Techinque_Stable_Diffusion_Video")
Interactive Demo 2: Stable Diffusion¶
In this demo, we will play with one of the most potent open-source diffusion models, Stable Diffusion 2.1, and try to connect with what we have learned.
Download the Stable Diffusion models¶
#@title Download the Stable Diffusion models
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler, PNDMScheduler
model_id = "stabilityai/stable-diffusion-2-1"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
# Use the PNDM scheduler as default
# pipe.scheduler = PNDMScheduler.from_config(pipe.scheduler.config)
# Use the DPMSolverMultistepScheduler (DPM-Solver++) scheduler here instead
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to(DEVICE)
Now you can let loose your imagination and create artworks from text!
Example prompts:
prompt = "A lovely cat running on the dessert in Van Gogh style, trending art."
prompt = "A ballerina dancing under the starry night in Monet style, trending art."
prompt = "A lovely cat running on the dessert in Van Gogh style, trending art." # @param {'type':'string'}
my_seed = 2023 # @param {'type':'integer'}
execute = False # @param {'type':'boolean'}
if execute:
image = pipe(prompt, num_inference_steps=50,
generator=torch.Generator("cuda").manual_seed(my_seed)).images[0]
image
Submit your feedback¶
# @title Submit your feedback
content_review(f"{feedback_prefix}_Stable_Diffusion_Interactive_Demo")
Think! 2: Architecture of Stable Diffusion Model¶
Can you see the similarity between the U-Net in Stable Diffusion and the baby UNet we defined up there?
To inspect the architecture, you can use the recursive_print(pipe.unet,deepest=2) function with a different deepest.
The text is encoded through the CLIP model, and you can also look at its structure below recursive_print(pipe.text_encoder,deepest=4), which is a large transformer!
Take 2 minutes to think and play with the code, then discuss as a group (~10 minutes).
Helper function to inspect network¶
# @title Helper function to inspect network
def recursive_print(module, prefix="", depth=0, deepest=3):
"""Simulating print(module) for torch.nn.Modules
but with depth control. Print to the `deepest` level. `deepest=0` means no print
"""
if depth == 0:
print(f"[{type(module).__name__}]")
if depth >= deepest:
return
for name, child in module.named_children():
if len([*child.named_children()]) == 0:
print(f"{prefix}({name}): {child}")
else:
if isinstance(child, nn.ModuleList):
print(f"{prefix}({name}): {type(child).__name__} len={len(child)}")
else:
print(f"{prefix}({name}): {type(child).__name__}")
recursive_print(child, prefix=prefix + " ", depth=depth + 1, deepest=deepest)
recursive_print(pipe.unet,deepest=2)
[UNet2DConditionModel]
(conv_in): Conv2d(4, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(time_proj): Timesteps()
(time_embedding): TimestepEmbedding
(linear_1): Linear(in_features=320, out_features=1280, bias=True)
(act): SiLU()
(linear_2): Linear(in_features=1280, out_features=1280, bias=True)
(down_blocks): ModuleList len=4
(0): CrossAttnDownBlock2D
(1): CrossAttnDownBlock2D
(2): CrossAttnDownBlock2D
(3): DownBlock2D
(up_blocks): ModuleList len=4
(0): UpBlock2D
(1): CrossAttnUpBlock2D
(2): CrossAttnUpBlock2D
(3): CrossAttnUpBlock2D
(mid_block): UNetMidBlock2DCrossAttn
(attentions): ModuleList len=1
(resnets): ModuleList len=2
(conv_norm_out): GroupNorm(32, 320, eps=1e-05, affine=True)
(conv_act): SiLU()
(conv_out): Conv2d(320, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
recursive_print(pipe.text_encoder,deepest=4)
[CLIPTextModel]
(text_model): CLIPTextTransformer
(embeddings): CLIPTextEmbeddings
(token_embedding): Embedding(49408, 1024)
(position_embedding): Embedding(77, 1024)
(encoder): CLIPEncoder
(layers): ModuleList len=23
(0): CLIPEncoderLayer
(1): CLIPEncoderLayer
(2): CLIPEncoderLayer
(3): CLIPEncoderLayer
(4): CLIPEncoderLayer
(5): CLIPEncoderLayer
(6): CLIPEncoderLayer
(7): CLIPEncoderLayer
(8): CLIPEncoderLayer
(9): CLIPEncoderLayer
(10): CLIPEncoderLayer
(11): CLIPEncoderLayer
(12): CLIPEncoderLayer
(13): CLIPEncoderLayer
(14): CLIPEncoderLayer
(15): CLIPEncoderLayer
(16): CLIPEncoderLayer
(17): CLIPEncoderLayer
(18): CLIPEncoderLayer
(19): CLIPEncoderLayer
(20): CLIPEncoderLayer
(21): CLIPEncoderLayer
(22): CLIPEncoderLayer
(final_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
Submit your feedback¶
# @title Submit your feedback
content_review(f"{feedback_prefix}_Architecture_of_Stable_Diffusion_Model_Discussion")
Section 3: Ethical Considerations¶
Video 4: Ethical Consideration¶
Submit your feedback¶
# @title Submit your feedback
content_review(f"{feedback_prefix}_Ethical_Consideration_Video")
Think! 3: Copyright of imagery generated from diffusion generated models¶
Suppose you prompt a pretrained diffusion model with the name of the artist and obtain beautiful imagery similar to that artist’s style. Who has the copyright of the generated image? The producing company of the diffusion model, the original artist, you, the prompter, the random seed and the weights, or the GPU that runs the inference?
Who do you think deserves the credit and why?
What if you apply enough post-processing steps to the generated images, e.g., finetune the prompt and seed, or edit the image?
Take 2 minutes to think in silence, then discuss as a group (~10 minutes).
Submit your feedback¶
# @title Submit your feedback
content_review(f"{feedback_prefix}_Copyrights_Discussion")
Summary¶
Today, we learned about
One major application for diffusion modeling, i.e., Modeling natural images.
Inductive biases suitable for image modeling: U-Net architecture and time modulation mechanism.
Introduction to conditional diffusion models, with a demo on Stable Diffusion.
Ethical considerations related to diffusion models, including copyright, misinformation, and fairness.
Daily survey¶
Don’t forget to complete your reflections and content check in the daily survey! Please be patient after logging in as there is a small delay before you will be redirected to the survey.
Bonus: Train Conditional Diffusion for MNIST¶
In this part, we’d like to train an MNIST generative model conditioned on the digit.
Here we will use a basic form of conditional modulation, i.e., digit embedding, to linearly control the relative gain of the features. After learning about the attention mechanism, you could think about better ways conditional modulation, e.g., using cross-attention to modulate the score model.
UNet score model with conditional modulation¶
class UNet_Conditional(nn.Module):
"""A time-dependent score-based model built upon U-Net architecture."""
def __init__(self, marginal_prob_std, channels=[32, 64, 128, 256], embed_dim=256,
text_dim=256, nClass=10):
"""Initialize a time-dependent score-based network.
Args:
marginal_prob_std: A function that takes time t and gives the standard
deviation of the perturbation kernel p_{0t}(x(t) | x(0)).
channels: The number of channels for feature maps of each resolution.
embed_dim: The dimensionality of Gaussian random feature embeddings of time.
text_dim: the embedding dimension of text / digits.
nClass: number of classes you want to model.
"""
super().__init__()
# random embedding for classes
self.cond_embed = nn.Embedding(nClass, text_dim)
# Gaussian random feature embedding layer for time
self.time_embed = nn.Sequential(
GaussianFourierProjection(embed_dim=embed_dim),
nn.Linear(embed_dim, embed_dim)
)
# Encoding layers where the resolution decreases
self.conv1 = nn.Conv2d(1, channels[0], 3, stride=1, bias=False)
self.t_mod1 = Dense(embed_dim, channels[0])
self.gnorm1 = nn.GroupNorm(4, num_channels=channels[0])
self.conv2 = nn.Conv2d(channels[0], channels[1], 3, stride=2, bias=False)
self.t_mod2 = Dense(embed_dim, channels[1])
self.gnorm2 = nn.GroupNorm(32, num_channels=channels[1])
self.y_mod2 = Dense(embed_dim, channels[1])
self.conv3 = nn.Conv2d(channels[1], channels[2], 3, stride=2, bias=False)
self.t_mod3 = Dense(embed_dim, channels[2])
self.gnorm3 = nn.GroupNorm(32, num_channels=channels[2])
self.y_mod3 = Dense(embed_dim, channels[2])
self.conv4 = nn.Conv2d(channels[2], channels[3], 3, stride=2, bias=False)
self.t_mod4 = Dense(embed_dim, channels[3])
self.gnorm4 = nn.GroupNorm(32, num_channels=channels[3])
self.y_mod4 = Dense(embed_dim, channels[3])
# Decoding layers where the resolution increases
self.tconv4 = nn.ConvTranspose2d(channels[3], channels[2], 3, stride=2, bias=False)
self.t_mod5 = Dense(embed_dim, channels[2])
self.y_mod5 = Dense(embed_dim, channels[2])
self.tgnorm4 = nn.GroupNorm(32, num_channels=channels[2])
self.tconv3 = nn.ConvTranspose2d(channels[2], channels[1], 3, stride=2, bias=False, output_padding=1) # + channels[2]
self.t_mod6 = Dense(embed_dim, channels[1])
self.y_mod6 = Dense(embed_dim, channels[1])
self.tgnorm3 = nn.GroupNorm(32, num_channels=channels[1])
self.tconv2 = nn.ConvTranspose2d(channels[1], channels[0], 3, stride=2, bias=False, output_padding=1) # + channels[1]
self.t_mod7 = Dense(embed_dim, channels[0])
self.y_mod7 = Dense(embed_dim, channels[0])
self.tgnorm2 = nn.GroupNorm(32, num_channels=channels[0])
self.tconv1 = nn.ConvTranspose2d(channels[0], 1, 3, stride=1)
# The swish activation function
self.act = nn.SiLU() # lambda x: x * torch.sigmoid(x)
self.marginal_prob_std = marginal_prob_std
for module in [self.y_mod2,self.y_mod3,self.y_mod4,
self.y_mod5,self.y_mod6,self.y_mod7]:
nn.init.normal_(module.dense.weight, mean=0, std=0.0001)
nn.init.constant_(module.dense.bias, 1.0)
def forward(self, x, t, y=None):
# Obtain the Gaussian random feature embedding for t
embed = self.act(self.time_embed(t))
y_embed = self.cond_embed(y)
# Encoding path
h1 = self.conv1(x) + self.t_mod1(embed)
## Incorporate information from t
## Group normalization
h1 = self.act(self.gnorm1(h1))
h2 = self.conv2(h1) + self.t_mod2(embed)
h2 = h2 * self.y_mod2(y_embed)
h2 = self.act(self.gnorm2(h2))
h3 = self.conv3(h2) + self.t_mod3(embed)
h3 = h3 * self.y_mod3(y_embed)
h3 = self.act(self.gnorm3(h3))
h4 = self.conv4(h3) + self.t_mod4(embed)
h4 = h4 * self.y_mod4(y_embed)
h4 = self.act(self.gnorm4(h4))
# Decoding path
h = self.tconv4(h4) + self.t_mod5(embed)
h = h * self.y_mod5(y_embed)
## Skip connection from the encoding path
h = self.act(self.tgnorm4(h))
h = self.tconv3(h + h3) + self.t_mod6(embed)
h = h * self.y_mod6(y_embed)
h = self.act(self.tgnorm3(h))
h = self.tconv2(h + h2) + self.t_mod7(embed)
h = h * self.y_mod7(y_embed)
h = self.act(self.tgnorm2(h))
h = self.tconv1(h + h1)
# Normalize output
h = h / self.marginal_prob_std(t)[:, None, None, None]
return h
Loss for conditional diffusion¶
def loss_fn_cond(model, x, y, marginal_prob_std, eps=1e-3):
"""The loss function for training score-based generative models.
Args:
model: A PyTorch model instance that represents a
time-dependent score-based model.
x: A mini-batch of training data.
marginal_prob_std: A function that gives the standard deviation of
the perturbation kernel.
eps: A tolerance value for numerical stability.
"""
random_t = torch.rand(x.shape[0], device=x.device) * (1. - eps) + eps
z = torch.randn_like(x)
std = marginal_prob_std(random_t)
perturbed_x = x + z * std[:, None, None, None]
score = model(perturbed_x, random_t, y=y)
loss = torch.mean(torch.sum((score * std[:, None, None, None] + z)**2,
dim=(1, 2, 3)))
return loss
Training conditional diffusion model¶
# @title Training conditional diffusion model
Lambda = 25 #@param {'type':'number'}
marginal_prob_std_fn = lambda t: marginal_prob_std(t, Lambda=Lambda, device=DEVICE)
diffusion_coeff_fn = lambda t: diffusion_coeff(t, Lambda=Lambda, device=DEVICE)
print("initilize new score model...")
score_model_cond = UNet_Conditional(marginal_prob_std=marginal_prob_std_fn)
score_model_cond = score_model_cond.to(DEVICE)
n_epochs = 10 # @param {'type':'integer'}
## size of a mini-batch
batch_size = 1024 # @param {'type':'integer'}
## learning rate
lr = 10e-4 # @param {'type':'number'}
dataset = MNIST('.', train=True, transform=transforms.ToTensor(), download=True)
data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=2)
optimizer = Adam(score_model_cond.parameters(), lr=lr)
scheduler = LambdaLR(optimizer, lr_lambda=lambda epoch: max(0.2, 0.99 ** epoch))
tqdm_epoch = trange(n_epochs)
for epoch in tqdm_epoch:
avg_loss = 0.
num_items = 0
for x, y in tqdm(data_loader):
x = x.to(DEVICE)
loss = loss_fn_cond(score_model_cond, x, y.to(DEVICE), marginal_prob_std_fn)
optimizer.zero_grad()
loss.backward()
optimizer.step()
avg_loss += loss.item() * x.shape[0]
num_items += x.shape[0]
scheduler.step()
lr_current = scheduler.get_last_lr()[0]
print('{} Average Loss: {:5f} lr {:.1e}'.format(epoch, avg_loss / num_items, lr_current))
# Print the averaged training loss so far.
tqdm_epoch.set_description('Average Loss: {:5f}'.format(avg_loss / num_items))
# Update the checkpoint after each epoch of training.
torch.save(score_model_cond.state_dict(), 'ckpt_cond.pth')
Note: The original value for n_epochs was 100.
Sample Conditional Diffusion¶
# @title Sample Conditional Diffusion
digit = 4 # @param {'type':'integer'}
sample_batch_size = 64 # @param {'type':'integer'}
num_steps = 250 # @param {'type':'integer'}
score_model_cond.eval()
## Generate samples using the specified sampler.
samples = Euler_Maruyama_sampler(
score_model_cond,
marginal_prob_std_fn,
diffusion_coeff_fn,
sample_batch_size,
num_steps=num_steps,
device=DEVICE,
y=digit*torch.ones(sample_batch_size, dtype=torch.long, device=DEVICE))
## Sample visualization.
samples = samples.clamp(0.0, 1.0)
%matplotlib inline
import matplotlib.pyplot as plt
sample_grid = make_grid(samples, nrow=int(np.sqrt(sample_batch_size)))
plt.figure(figsize=(6, 6))
plt.axis('off')
plt.imshow(sample_grid.permute(1, 2, 0).cpu(), vmin=0., vmax=1.)
plt.show()
torch.cuda.empty_cache()