# -*- coding: utf-8 -*-
"""01_model.ipynb
Automatically generated by Colaboratory.
Original file is located at
https://colab.research.google.com/drive/1OWXPL8K-jKC4KgGmYXXkeBWOL2V1biuf
# Setup
"""
import torch
from fastai.callbacks import *
from fastai.vision import *
H = 160
W = 192
D = 128
def conv_block(c_in, c_out, ks, num_groups=None, **conv_kwargs):
"A sequence of modules composed of Group Norm, ReLU and Conv3d in order"
if not num_groups : num_groups = int(c_in/2) if c_in%2 == 0 else None
return nn.Sequential(nn.GroupNorm(num_groups, c_in),
nn.ReLU(),
nn.Conv3d(c_in, c_out, ks, **conv_kwargs))
def reslike_block(nf, num_groups=None, bottle_neck:bool=False, **conv_kwargs):
"A ResNet-like block with the GroupNorm normalization providing optional bottle-neck functionality"
nf_inner = nf / 2 if bottle_neck else nf
return SequentialEx(conv_block(num_groups=num_groups, c_in=nf, c_out=nf_inner, ks=3, stride=1, padding=1, **conv_kwargs),
conv_block(num_groups=num_groups, c_in=nf_inner, c_out=nf, ks=3, stride=1, padding=1, **conv_kwargs),
MergeLayer())
def upsize(c_in, c_out, ks=1, scale=2):
"Reduce the number of features by 2 using Conv with kernel size 1x1x1 and double the spatial dimension using 3D trilinear upsampling"
return nn.Sequential(nn.Conv3d(c_in, c_out, ks),
nn.Upsample(scale_factor=scale, mode='trilinear', align_corners=True))
def hook_debug(module, input, output):
"""
Print out what's been hooked usually for debugging purpose
----------------------------------------------------------
Example:
Hooks(ms, hook_debug, is_forward=True, detach=False)
"""
print('Hooking ' + module.__class__.__name__)
print('output size:', output.data.size())
return output
class Encoder(nn.Module):
"Encoder part"
def __init__(self):
super().__init__()
self.conv1 = nn.Conv3d(4, 32, 3, stride=1, padding=1)
self.res_block1 = reslike_block(32, num_groups=8)
self.conv_block1 = conv_block(32, 64, 3, num_groups=8, stride=2, padding=1)
self.res_block2 = reslike_block(64, num_groups=8)
self.conv_block2 = conv_block(64, 64, 3, num_groups=8, stride=1, padding=1)
self.res_block3 = reslike_block(64, num_groups=8)
self.conv_block3 = conv_block(64, 128, 3, num_groups=8, stride=2, padding=1)
self.res_block4 = reslike_block(128, num_groups=8)
self.conv_block4 = conv_block(128, 128, 3, num_groups=8, stride=1, padding=1)
self.res_block5 = reslike_block(128, num_groups=8)
self.conv_block5 = conv_block(128, 256, 3, num_groups=8, stride=2, padding=1)
self.res_block6 = reslike_block(256, num_groups=8)
self.conv_block6 = conv_block(256, 256, 3, num_groups=8, stride=1, padding=1)
self.res_block7 = reslike_block(256, num_groups=8)
self.conv_block7 = conv_block(256, 256, 3, num_groups=8, stride=1, padding=1)
self.res_block8 = reslike_block(256, num_groups=8)
self.conv_block8 = conv_block(256, 256, 3, num_groups=8, stride=1, padding=1)
self.res_block9 = reslike_block(256, num_groups=8)
def forward(self, x):
x = self.conv1(x) # Output size: (1, 32, 160, 192, 128)
x = self.res_block1(x) # Output size: (1, 32, 160, 192, 128)
x = self.conv_block1(x) # Output size: (1, 64, 80, 96, 64)
x = self.res_block2(x) # Output size: (1, 64, 80, 96, 64)
x = self.conv_block2(x) # Output size: (1, 64, 80, 96, 64)
x = self.res_block3(x) # Output size: (1, 64, 80, 96, 64)
x = self.conv_block3(x) # Output size: (1, 128, 40, 48, 32)
x = self.res_block4(x) # Output size: (1, 128, 40, 48, 32)
x = self.conv_block4(x) # Output size: (1, 128, 40, 48, 32)
x = self.res_block5(x) # Output size: (1, 128, 40, 48, 32)
x = self.conv_block5(x) # Output size: (1, 256, 20, 24, 16)
x = self.res_block6(x) # Output size: (1, 256, 20, 24, 16)
x = self.conv_block6(x) # Output size: (1, 256, 20, 24, 16)
x = self.res_block7(x) # Output size: (1, 256, 20, 24, 16)
x = self.conv_block7(x) # Output size: (1, 256, 20, 24, 16)
x = self.res_block8(x) # Output size: (1, 256, 20, 24, 16)
x = self.conv_block8(x) # Output size: (1, 256, 20, 24, 16)
x = self.res_block9(x) # Output size: (1, 256, 20, 24, 16)
return x
class Decoder(nn.Module):
"Decoder Part"
def __init__(self):
super().__init__()
self.upsize1 = upsize(256, 128)
self.reslike1 = reslike_block(128, num_groups=8)
self.upsize2 = upsize(128, 64)
self.reslike2 = reslike_block(64, num_groups=8)
self.upsize3 = upsize(64, 32)
self.reslike3 = reslike_block(32, num_groups=8)
self.conv1 = nn.Conv3d(32, 3, 1)
self.sigmoid1 = torch.nn.Sigmoid()
def forward(self, x):
x = self.upsize1(x) # Output size: (1, 128, 40, 48, 32)
x = x + hooks.stored[2] # Output size: (1, 128, 40, 48, 32)
x = self.reslike1(x) # Output size: (1, 128, 40, 48, 32)
x = self.upsize2(x) # Output size: (1, 64, 80, 96, 64)
x = x + hooks.stored[1] # Output size: (1, 64, 80, 96, 64)
x = self.reslike2(x) # Output size: (1, 64, 80, 96, 64)
x = self.upsize3(x) # Output size: (1, 32, 160, 192, 128)
x = x + hooks.stored[0] # Output size: (1, 32, 160, 192, 128)
x = self.reslike3(x) # Output size: (1, 32, 160, 192, 128)
x = self.conv1(x) # Output size: (1, 3, 160, 192, 128)
x = self.sigmoid1(x) # Output size: (1, 3, 160, 192, 128)
return x
class VAEEncoder(nn.Module):
"Variational auto-encoder encoder part"
def __init__(self, latent_dim:int=128):
super().__init__()
self.latent_dim = latent_dim
self.conv_block = conv_block(256, 16, 3, num_groups=8, stride=2, padding=1)
self.linear1 = nn.Linear(60, 1)
# Assumed latent variable's probability density function parameters
self.z_mean = nn.Linear(256, latent_dim)
self.z_log_var = nn.Linear(256, latent_dim)
self.epsilon = nn.Parameter(torch.randn(1, latent_dim))
def forward(self, x):
x = self.conv_block(x) # Output size: (1, 16, 10, 12, 8)
x = x.view(256, -1) # Output size: (256, 60)
x = self.linear1(x) # Output size: (256, 1)
x = x.view(1, 256) # Output size: (1, 256)
z_mean = self.z_mean(x) # Output size: (1, 128)
z_var = self.z_log_var(x).exp() # Output size: (1, 128)
return z_mean + z_var * self.epsilon # Output size: (1, 128)
class VAEDecoder(nn.Module):
"Variational auto-encoder decoder part"
def __init__(self):
super().__init__()
self.linear1 = nn.Linear(128, 256*60)
self.relu1 = nn.ReLU()
self.upsize1 = upsize(16, 256)
self.upsize2 = upsize(256, 128)
self.reslike1 = reslike_block(128, num_groups=8)
self.upsize3 = upsize(128, 64)
self.reslike2 = reslike_block(64, num_groups=8)
self.upsize4 = upsize(64, 32)
self.reslike3 = reslike_block(32, num_groups=8)
self.conv1 = nn.Conv3d(32, 4, 1)
def forward(self, x):
x = self.linear1(x) # Output size: (1, 256*60)
x = self.relu1(x) # Output size: (1, 256*60)
x = x.view(1, 16, 10, 12, 8) # Output size: (1, 16, 10, 12, 8)
x = self.upsize1(x) # Output size: (1, 256, 20, 24, 16)
x = self.upsize2(x) # Output size: (1, 128, 40, 48, 32)
x = self.reslike1(x) # Output size: (1, 128, 40, 48, 32)
x = self.upsize3(x) # Output size: (1, 64, 80, 96, 64)
x = self.reslike2(x) # Output size: (1, 64, 80, 96, 64)
x = self.upsize4(x) # Output size: (1, 32, 160, 192, 128)
x = self.reslike3(x) # Output size: (1, 32, 160, 192, 128)
x = self.conv1(x) # Output size: (1, 4, 160, 192, 128)
return x
class AutoUNet(nn.Module):
"3D U-Net using autoencoder regularization"
def __init__(self):
super().__init__()
self.encoder = Encoder()
self.decoder = Decoder()
self.vencoder = VAEEncoder(latent_dim=128)
self.vdecoder = VAEDecoder()
def forward(self, input):
interm_res = self.encoder(input)
top_res = self.decoder(interm_res) # Output size: (1, 3, 160, 192, 128)
bottom_res = self.vdecoder(self.vencoder(interm_res)) # Output size: (1, 4, 160, 192, 128)
return top_res, bottom_res
class SoftDiceLoss(nn.Module):
"Soft dice loss based on a measure of overlap between prediction and ground truth"
def __init__(self, epsilon=1e-6, c=3):
super().__init__()
self.epsilon = epsilon
self.c = 3
def forward(self, x:Tensor, y:Tensor):
intersection = 2 * ( (x*y).sum() )
union = (x**2).sum() + (y**2).sum()
return 1 - ( ( intersection / (union + self.epsilon) ) / self.c )
class KLDivergence(nn.Module):
"KL divergence between the estimated normal distribution and a prior distribution"
N = H * W * D #hyperparameter check
def __init__(self):
super().__init__()
def forward(self, z_mean:Tensor, z_log_var:Tensor):
z_var = z_log_var.exp()
return (1/self.N) * ( (z_mean**2 + z_var**2 - z_log_var**2 - 1).sum() )
class L2Loss(nn.Module):
"Measuring the `Euclidian distance` between prediction and ground truh using `L2 Norm`"
def __init__(self):
super().__init__()
def forward(self, x:Tensor, y:Tensor):
return ( (x - y)**2 ).sum()
autounet = AutoUNet()
ms = [autounet.encoder.res_block1,
autounet.encoder.res_block3,
autounet.encoder.res_block5,
autounet.vencoder.z_mean,
autounet.vencoder.z_log_var]
hooks = hook_outputs(ms, detach=False, grad=False)
lr = 1e-4
optimizer = optim.Adam(autounet.parameters(), lr)
