Denoising
Teaching a network to repair.
Take a clean photo, and scratch it. Now you've got a perfect training example. The scratched photo is the input, and the clean one is the answer. Since you made the damage yourself, you can make unlimited training data. A network trained like this learns what real images should look like.
- Damage a clean photo
- Input: damaged. Answer: clean
- Unlimited training data
Removing noise
Just a small change. Before each step, we add random noise to the images. The input is noisy, but the target is the clean image. Run it. The top row is heavily noisy digits. The bottom row is the network's repair, and they're clean and clear. Think about what it's learned here. It has an internal idea of what digits look like, so it can pull a good image out of the noise. Keep that thought, because in a few lessons, it's going to be the key to everything.
import torch
from torch import nn
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
root = "data"
train = datasets.MNIST(root, train=True, transform=transforms.ToTensor())
clean = datasets.MNIST(root, train=False).data[:6].float().unsqueeze(1) / 255
torch.manual_seed(0)
model = nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.ReLU(),
nn.Linear(256, 784), nn.Sigmoid(), nn.Unflatten(1, (1, 28, 28)))
opt = torch.optim.Adam(model.parameters(), lr=2e-3)
dl = torch.utils.data.DataLoader(train, batch_size=256, shuffle=True)
for epoch in range(4):
for x, _ in dl:
noisy = (x + 0.5 * torch.randn_like(x)).clamp(0, 1)
loss = ((model(noisy) - x) ** 2).mean()
opt.zero_grad(); loss.backward(); opt.step()
noisy = (clean + 0.5 * torch.randn_like(clean)).clamp(0, 1)
with torch.no_grad():
fixed = model(noisy)
fig, ax = plt.subplots(2, 6, figsize=(9, 3.2))
for i in range(6):
ax[0, i].imshow(noisy[i, 0], cmap="gray"); ax[1, i].imshow(fixed[i, 0], cmap="gray")
ax[0, i].axis("off"); ax[1, i].axis("off")
plt.tight_layout()
plt.savefig("denoise.png", dpi=120)
print("Top: noisy input. Bottom: the network's repair.")Top: noisy input. Bottom: the network's repair.
Summary
A denoiser learns what real data looks like, so it can recover images from noise. Hold that idea. But first, there's a simpler question. Could the decoder of an autoencoder create completely new images?
Key takeaways- Add noise, target the clean image
- The network learns what images look like
- Keep this idea in mind