Tutorial 2: Image Diffusion#
Week 2, Day 4: Diffusion Generative Models
By Neuromatch Academy
Content creators: Binxu Wang
Content reviewers: Shaonan Wang, Dongrui Deng, Dora Zhiyu Yang, Adrita Das, Jiaxin Cindy Tu
Content editors: Shaonan Wang, Jiaxin Cindy Tu
Production editors: Spiros Chavlis, Konstantine Tsafatinos
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#
Install and import feedback gadget#
Show code cell source
# @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_T2"
Uncomment to install dependencies#
Show code cell source
# @title Uncomment to install dependencies
# !pip install diffusers transformers tokenizers --quiet
# !pip install accelerate --quiet
# 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 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#
Show code cell source
# @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
Show code cell source
# @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.
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()#
Show code cell source
# @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.
Section 1: 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#
Show code cell source
# @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 \(\mathscr 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'):
r"""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)
#################################################
raise NotImplementedError("Student exercise: Implement the standard deviation")
#################################################
std = ...
return std
def diffusion_coeff(t, Lambda, device='cpu'):
r"""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.
"""
#################################################
raise NotImplementedError("Student exercise: Implement the diffusion coefficients")
#################################################
diff_coeff = ...
return diff_coeff.to(device)
Submit your feedback#
Show code cell source
# @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#
Show code cell source
# @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#
Show code cell source
# @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#
Show code cell source
# @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'):
r"""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)
#################################################
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#
Show code cell source
# @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#
Show code cell source
# @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#
Show code cell source
# @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')
Define the Sampler#
Show code cell source
# @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#
Show code cell source
# @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[21], 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/hostedtoolcache/Python/3.10.20/x64/lib/python3.10/site-packages/torch/serialization.py:1484, in load(f, map_location, pickle_module, weights_only, mmap, **pickle_load_args)
1481 if "encoding" not in pickle_load_args.keys():
1482 pickle_load_args["encoding"] = "utf-8"
-> 1484 with _open_file_like(f, "rb") as opened_file:
1485 if _is_zipfile(opened_file):
1486 # The zipfile reader is going to advance the current file position.
1487 # If we want to actually tail call to torch.jit.load, we need to
1488 # reset back to the original position.
1489 orig_position = opened_file.tell()
File /opt/hostedtoolcache/Python/3.10.20/x64/lib/python3.10/site-packages/torch/serialization.py:759, in _open_file_like(name_or_buffer, mode)
757 def _open_file_like(name_or_buffer: FileLike, mode: str) -> _opener[IO[bytes]]:
758 if _is_path(name_or_buffer):
--> 759 return _open_file(name_or_buffer, mode)
760 else:
761 if "w" in mode:
File /opt/hostedtoolcache/Python/3.10.20/x64/lib/python3.10/site-packages/torch/serialization.py:740, in _open_file.__init__(self, name, mode)
739 def __init__(self, name: Union[str, os.PathLike[str]], mode: str) -> None:
--> 740 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: Ethical Considerations#
Video 2: Ethical Consideration#
Submit your feedback#
Show code cell source
# @title Submit your feedback
content_review(f"{feedback_prefix}_Ethical_Consideration_Video")
Think! 2: 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#
Show code cell source
# @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.
Ethical considerations related to diffusion models, including copyright, misinformation, and fairness.