chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,187 @@
|
||||
"""
|
||||
---
|
||||
title: U-Net
|
||||
summary: >
|
||||
PyTorch implementation and tutorial of U-Net model.
|
||||
---
|
||||
|
||||
# U-Net
|
||||
|
||||
This is an implementation of the U-Net model from the paper,
|
||||
[U-Net: Convolutional Networks for Biomedical Image Segmentation](https://arxiv.org/abs/1505.04597).
|
||||
|
||||
U-Net consists of a contracting path and an expansive path.
|
||||
The contracting path is a series of convolutional layers and pooling layers,
|
||||
where the resolution of the feature map gets progressively reduced.
|
||||
Expansive path is a series of up-sampling layers and convolutional layers
|
||||
where the resolution of the feature map gets progressively increased.
|
||||
|
||||
At every step in the expansive path the corresponding feature map from the contracting path
|
||||
concatenated with the current feature map.
|
||||
|
||||
|
||||
|
||||
Here is the [training code](experiment.html) for an experiment that trains a U-Net
|
||||
on [Carvana dataset](carvana.html).
|
||||
"""
|
||||
import torch
|
||||
import torchvision.transforms.functional
|
||||
from torch import nn
|
||||
|
||||
|
||||
class DoubleConvolution(nn.Module):
|
||||
"""
|
||||
### Two $3 \times 3$ Convolution Layers
|
||||
|
||||
Each step in the contraction path and expansive path have two $3 \times 3$
|
||||
convolutional layers followed by ReLU activations.
|
||||
|
||||
In the U-Net paper they used $0$ padding,
|
||||
but we use $1$ padding so that final feature map is not cropped.
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels: int, out_channels: int):
|
||||
"""
|
||||
:param in_channels: is the number of input channels
|
||||
:param out_channels: is the number of output channels
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# First $3 \times 3$ convolutional layer
|
||||
self.first = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
|
||||
self.act1 = nn.ReLU()
|
||||
# Second $3 \times 3$ convolutional layer
|
||||
self.second = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
|
||||
self.act2 = nn.ReLU()
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
# Apply the two convolution layers and activations
|
||||
x = self.first(x)
|
||||
x = self.act1(x)
|
||||
x = self.second(x)
|
||||
return self.act2(x)
|
||||
|
||||
|
||||
class DownSample(nn.Module):
|
||||
"""
|
||||
### Down-sample
|
||||
|
||||
Each step in the contracting path down-samples the feature map with
|
||||
a $2 \times 2$ max pooling layer.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# Max pooling layer
|
||||
self.pool = nn.MaxPool2d(2)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
return self.pool(x)
|
||||
|
||||
|
||||
class UpSample(nn.Module):
|
||||
"""
|
||||
### Up-sample
|
||||
|
||||
Each step in the expansive path up-samples the feature map with
|
||||
a $2 \times 2$ up-convolution.
|
||||
"""
|
||||
def __init__(self, in_channels: int, out_channels: int):
|
||||
super().__init__()
|
||||
|
||||
# Up-convolution
|
||||
self.up = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
return self.up(x)
|
||||
|
||||
|
||||
class CropAndConcat(nn.Module):
|
||||
"""
|
||||
### Crop and Concatenate the feature map
|
||||
|
||||
At every step in the expansive path the corresponding feature map from the contracting path
|
||||
concatenated with the current feature map.
|
||||
"""
|
||||
def forward(self, x: torch.Tensor, contracting_x: torch.Tensor):
|
||||
"""
|
||||
:param x: current feature map in the expansive path
|
||||
:param contracting_x: corresponding feature map from the contracting path
|
||||
"""
|
||||
|
||||
# Crop the feature map from the contracting path to the size of the current feature map
|
||||
contracting_x = torchvision.transforms.functional.center_crop(contracting_x, [x.shape[2], x.shape[3]])
|
||||
# Concatenate the feature maps
|
||||
x = torch.cat([x, contracting_x], dim=1)
|
||||
#
|
||||
return x
|
||||
|
||||
|
||||
class UNet(nn.Module):
|
||||
"""
|
||||
## U-Net
|
||||
"""
|
||||
def __init__(self, in_channels: int, out_channels: int):
|
||||
"""
|
||||
:param in_channels: number of channels in the input image
|
||||
:param out_channels: number of channels in the result feature map
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Double convolution layers for the contracting path.
|
||||
# The number of features gets doubled at each step starting from $64$.
|
||||
self.down_conv = nn.ModuleList([DoubleConvolution(i, o) for i, o in
|
||||
[(in_channels, 64), (64, 128), (128, 256), (256, 512)]])
|
||||
# Down sampling layers for the contracting path
|
||||
self.down_sample = nn.ModuleList([DownSample() for _ in range(4)])
|
||||
|
||||
# The two convolution layers at the lowest resolution (the bottom of the U).
|
||||
self.middle_conv = DoubleConvolution(512, 1024)
|
||||
|
||||
# Up sampling layers for the expansive path.
|
||||
# The number of features is halved with up-sampling.
|
||||
self.up_sample = nn.ModuleList([UpSample(i, o) for i, o in
|
||||
[(1024, 512), (512, 256), (256, 128), (128, 64)]])
|
||||
# Double convolution layers for the expansive path.
|
||||
# Their input is the concatenation of the current feature map and the feature map from the
|
||||
# contracting path. Therefore, the number of input features is double the number of features
|
||||
# from up-sampling.
|
||||
self.up_conv = nn.ModuleList([DoubleConvolution(i, o) for i, o in
|
||||
[(1024, 512), (512, 256), (256, 128), (128, 64)]])
|
||||
# Crop and concatenate layers for the expansive path.
|
||||
self.concat = nn.ModuleList([CropAndConcat() for _ in range(4)])
|
||||
# Final $1 \times 1$ convolution layer to produce the output
|
||||
self.final_conv = nn.Conv2d(64, out_channels, kernel_size=1)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
:param x: input image
|
||||
"""
|
||||
# To collect the outputs of contracting path for later concatenation with the expansive path.
|
||||
pass_through = []
|
||||
# Contracting path
|
||||
for i in range(len(self.down_conv)):
|
||||
# Two $3 \times 3$ convolutional layers
|
||||
x = self.down_conv[i](x)
|
||||
# Collect the output
|
||||
pass_through.append(x)
|
||||
# Down-sample
|
||||
x = self.down_sample[i](x)
|
||||
|
||||
# Two $3 \times 3$ convolutional layers at the bottom of the U-Net
|
||||
x = self.middle_conv(x)
|
||||
|
||||
# Expansive path
|
||||
for i in range(len(self.up_conv)):
|
||||
# Up-sample
|
||||
x = self.up_sample[i](x)
|
||||
# Concatenate the output of the contracting path
|
||||
x = self.concat[i](x, pass_through.pop())
|
||||
# Two $3 \times 3$ convolutional layers
|
||||
x = self.up_conv[i](x)
|
||||
|
||||
# Final $1 \times 1$ convolution layer
|
||||
x = self.final_conv(x)
|
||||
|
||||
#
|
||||
return x
|
||||
@@ -0,0 +1,82 @@
|
||||
"""
|
||||
---
|
||||
title: Carvana dataset for the U-Net experiment
|
||||
summary: >
|
||||
Carvana dataset for the U-Net experiment.
|
||||
---
|
||||
|
||||
# Carvana Dataset for the [U-Net](index.html) [experiment](experiment.html)
|
||||
|
||||
You can find the download instructions
|
||||
[on Kaggle](https://www.kaggle.com/competitions/carvana-image-masking-challenge/data).
|
||||
|
||||
Save the training images inside `carvana/train` folder and the masks in `carvana/train_masks` folder.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import torchvision.transforms.functional
|
||||
from PIL import Image
|
||||
|
||||
import torch.utils.data
|
||||
from labml import lab
|
||||
|
||||
|
||||
class CarvanaDataset(torch.utils.data.Dataset):
|
||||
"""
|
||||
## Carvana Dataset
|
||||
"""
|
||||
|
||||
def __init__(self, image_path: Path, mask_path: Path):
|
||||
"""
|
||||
:param image_path: is the path to the images
|
||||
:param mask_path: is the path to the masks
|
||||
"""
|
||||
# Get a dictionary of images by id
|
||||
self.images = {p.stem: p for p in image_path.iterdir()}
|
||||
# Get a dictionary of masks by id
|
||||
self.masks = {p.stem[:-5]: p for p in mask_path.iterdir()}
|
||||
|
||||
# Image ids list
|
||||
self.ids = list(self.images.keys())
|
||||
|
||||
# Transformations
|
||||
self.transforms = torchvision.transforms.Compose([
|
||||
torchvision.transforms.Resize(572),
|
||||
torchvision.transforms.ToTensor(),
|
||||
])
|
||||
|
||||
def __getitem__(self, idx: int):
|
||||
"""
|
||||
#### Get an image and its mask.
|
||||
|
||||
:param idx: is index of the image
|
||||
"""
|
||||
|
||||
# Get image id
|
||||
id_ = self.ids[idx]
|
||||
# Load image
|
||||
image = Image.open(self.images[id_])
|
||||
# Transform image and convert it to a PyTorch tensor
|
||||
image = self.transforms(image)
|
||||
# Load mask
|
||||
mask = Image.open(self.masks[id_])
|
||||
# Transform mask and convert it to a PyTorch tensor
|
||||
mask = self.transforms(mask)
|
||||
|
||||
# The mask values were not $1$, so we scale it appropriately.
|
||||
mask = mask / mask.max()
|
||||
|
||||
# Return the image and the mask
|
||||
return image, mask
|
||||
|
||||
def __len__(self):
|
||||
"""
|
||||
#### Size of the dataset
|
||||
"""
|
||||
return len(self.ids)
|
||||
|
||||
|
||||
# Testing code
|
||||
if __name__ == '__main__':
|
||||
ds = CarvanaDataset(lab.get_data_path() / 'carvana' / 'train', lab.get_data_path() / 'carvana' / 'train_masks')
|
||||
@@ -0,0 +1,169 @@
|
||||
"""
|
||||
---
|
||||
title: Training a U-Net on Carvana dataset
|
||||
summary: >
|
||||
Code for training a U-Net model on Carvana dataset.
|
||||
---
|
||||
|
||||
# Training [U-Net](index.html)
|
||||
|
||||
This trains a [U-Net](index.html) model on [Carvana dataset](carvana.html).
|
||||
You can find the download instructions
|
||||
[on Kaggle](https://www.kaggle.com/competitions/carvana-image-masking-challenge/data).
|
||||
|
||||
Save the training images inside `carvana/train` folder and the masks in `carvana/train_masks` folder.
|
||||
|
||||
For simplicity, we do not do a training and validation split.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import torchvision.transforms.functional
|
||||
|
||||
import torch
|
||||
import torch.utils.data
|
||||
from labml import lab, tracker, experiment, monit
|
||||
from labml.configs import BaseConfigs
|
||||
from labml_nn.helpers.device import DeviceConfigs
|
||||
from labml_nn.unet import UNet
|
||||
from labml_nn.unet.carvana import CarvanaDataset
|
||||
from torch import nn
|
||||
|
||||
|
||||
class Configs(BaseConfigs):
|
||||
"""
|
||||
## Configurations
|
||||
"""
|
||||
# Device to train the model on.
|
||||
# [`DeviceConfigs`](../helpers/device.html)
|
||||
# picks up an available CUDA device or defaults to CPU.
|
||||
device: torch.device = DeviceConfigs()
|
||||
|
||||
# [U-Net](index.html) model
|
||||
model: UNet
|
||||
|
||||
# Number of channels in the image. $3$ for RGB.
|
||||
image_channels: int = 3
|
||||
# Number of channels in the output mask. $1$ for binary mask.
|
||||
mask_channels: int = 1
|
||||
|
||||
# Batch size
|
||||
batch_size: int = 1
|
||||
# Learning rate
|
||||
learning_rate: float = 2.5e-4
|
||||
|
||||
# Number of training epochs
|
||||
epochs: int = 4
|
||||
|
||||
# Dataset
|
||||
dataset: CarvanaDataset
|
||||
# Dataloader
|
||||
data_loader: torch.utils.data.DataLoader
|
||||
|
||||
# Loss function
|
||||
loss_func = nn.BCELoss()
|
||||
# Sigmoid function for binary classification
|
||||
sigmoid = nn.Sigmoid()
|
||||
|
||||
# Adam optimizer
|
||||
optimizer: torch.optim.Adam
|
||||
|
||||
def init(self):
|
||||
# Initialize the [Carvana dataset](carvana.html)
|
||||
self.dataset = CarvanaDataset(lab.get_data_path() / 'carvana' / 'train',
|
||||
lab.get_data_path() / 'carvana' / 'train_masks')
|
||||
# Initialize the model
|
||||
self.model = UNet(self.image_channels, self.mask_channels).to(self.device)
|
||||
|
||||
# Create dataloader
|
||||
self.data_loader = torch.utils.data.DataLoader(self.dataset, self.batch_size,
|
||||
shuffle=True, pin_memory=True)
|
||||
# Create optimizer
|
||||
self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate)
|
||||
|
||||
# Image logging
|
||||
tracker.set_image("sample", True)
|
||||
|
||||
@torch.no_grad()
|
||||
def sample(self, idx=-1):
|
||||
"""
|
||||
### Sample images
|
||||
"""
|
||||
|
||||
# Get a random sample
|
||||
x, _ = self.dataset[np.random.randint(len(self.dataset))]
|
||||
# Move data to device
|
||||
x = x.to(self.device)
|
||||
|
||||
# Get predicted mask
|
||||
mask = self.sigmoid(self.model(x[None, :]))
|
||||
# Crop the image to the size of the mask
|
||||
x = torchvision.transforms.functional.center_crop(x, [mask.shape[2], mask.shape[3]])
|
||||
# Log samples
|
||||
tracker.save('sample', x * mask)
|
||||
|
||||
def train(self):
|
||||
"""
|
||||
### Train for an epoch
|
||||
"""
|
||||
|
||||
# Iterate through the dataset.
|
||||
# Use [`mix`](https://docs.labml.ai/api/monit.html#labml.monit.mix)
|
||||
# to sample $50$ times per epoch.
|
||||
for _, (image, mask) in monit.mix(('Train', self.data_loader), (self.sample, list(range(50)))):
|
||||
# Increment global step
|
||||
tracker.add_global_step()
|
||||
# Move data to device
|
||||
image, mask = image.to(self.device), mask.to(self.device)
|
||||
|
||||
# Make the gradients zero
|
||||
self.optimizer.zero_grad()
|
||||
# Get predicted mask logits
|
||||
logits = self.model(image)
|
||||
# Crop the target mask to the size of the logits. Size of the logits will be smaller if we
|
||||
# don't use padding in convolutional layers in the U-Net.
|
||||
mask = torchvision.transforms.functional.center_crop(mask, [logits.shape[2], logits.shape[3]])
|
||||
# Calculate loss
|
||||
loss = self.loss_func(self.sigmoid(logits), mask)
|
||||
# Compute gradients
|
||||
loss.backward()
|
||||
# Take an optimization step
|
||||
self.optimizer.step()
|
||||
# Track the loss
|
||||
tracker.save('loss', loss)
|
||||
|
||||
def run(self):
|
||||
"""
|
||||
### Training loop
|
||||
"""
|
||||
for _ in monit.loop(self.epochs):
|
||||
# Train the model
|
||||
self.train()
|
||||
# New line in the console
|
||||
tracker.new_line()
|
||||
# Save the model
|
||||
|
||||
|
||||
def main():
|
||||
# Create experiment
|
||||
experiment.create(name='unet')
|
||||
|
||||
# Create configurations
|
||||
configs = Configs()
|
||||
|
||||
# Set configurations. You can override the defaults by passing the values in the dictionary.
|
||||
experiment.configs(configs, {})
|
||||
|
||||
# Initialize
|
||||
configs.init()
|
||||
|
||||
# Set models for saving and loading
|
||||
experiment.add_pytorch_models({'model': configs.model})
|
||||
|
||||
# Start and run the training loop
|
||||
with experiment.start():
|
||||
configs.run()
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user