Skip to content

Instantly share code, notes, and snippets.

Created July 9, 2020 21:28
  • Star 2 You must be signed in to star a gist
  • Fork 1 You must be signed in to fork a gist
Star You must be signed in to star a gist
What would you like to do?
U-Net in PyTorch
U-Net architecture in PyTorch (
Author: Jacob Reinhold (
import torch
from torch import nn
from torch.nn import functional as F
class ConvLayer(nn.Sequential):
def __init__(self, in_channels:int, out_channels:int):
self.add_module('conv', nn.Conv2d(in_channels, out_channels,
3, padding=1, bias=False))
self.add_module('norm', nn.BatchNorm2d(out_channels))
self.add_module('relu', nn.ReLU(inplace=True))
class UNetBlock(nn.Sequential):
def __init__(self, in_channels:int, out_channels:int):
self.add_module('block1', ConvLayer(in_channels, out_channels))
self.add_module('block2', ConvLayer(out_channels, out_channels))
class UNet(nn.Module):
def __init__(self, in_channels:int, out_channels:int, channel_base:int=64):
self.down_layers = nn.ModuleList([])
n_chan = lambda x: channel_base*2**x
self.down_layers.append(UNetBlock(in_channels, n_chan(0)))
for i in range(3):
self.down_layers.append(UNetBlock(n_chan(i), n_chan(i+1)))
self.bottleneck = UNetBlock(n_chan(3), n_chan(4))
self.up_layers = nn.ModuleList([])
for i in reversed(range(1, 4)):
self.up_layers.append(UNetBlock(n_chan(i+1)+n_chan(i), n_chan(i)))
UNetBlock(n_chan(1)+n_chan(0), n_chan(0),),
nn.Conv2d(n_chan(0), out_channels, 1)))
def interp_cat(x, skip):
x = F.interpolate(x, skip.shape[2:], mode='bilinear', align_corners=True)
return, skip), 1)
def forward(self, x):
skip_connections = []
for down_layer in self.down_layers:
x = down_layer(x)
x = F.max_pool2d(x, 2)
x = self.bottleneck(x)
for up_layer in self.up_layers:
skip = skip_connections.pop()
x = self.interp_cat(x, skip)
x = up_layer(x)
return x
if __name__ == "__main__":
model = UNet(1,1)
x = torch.randn(1,1,128,128)
Copy link

Implementation of a 2D U-Net in PyTorch. Differences from original: 1) uses linear interpolation instead of transposed conv. as upsampling, 2) maintains the input size by padding. Not tested extensively.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment