Here is the code, which was originally written for 2D image registration,(recursive cascade network) I modified for 3D image. But not being able to match the size of the layer, problem is occurring in concatenation operations.
import torch
import torch.nn as nn
from torch.nn import ReLU, LeakyReLU
import torch.nn.functional as F
def conv(dim = 3):
return nn.Conv3d
def trans_conv(dim = 3):
return nn.ConvTranspose3d
def convolve(in_channels, out_channels, kernel_size, stride, dim=3):
return conv(dim=dim)(in_channels, out_channels, kernel_size, stride=stride, padding=1)
def convolveReLU(in_channels, out_channels, kernel_size, stride, dim=3):
return nn.Sequential(ReLU, convolve(in_channels, out_channels, kernel_size, stride, dim=dim))
def convolveLeakyReLU(in_channels, out_channels, kernel_size, stride, dim=3):
return nn.Sequential(LeakyReLU(0.1), convolve(in_channels, out_channels, kernel_size, stride, dim=dim))
def upconvolve(in_channels, out_channels, kernel_size, stride, dim=3):
return trans_conv(dim=dim)(in_channels, out_channels, kernel_size, stride, padding=1)
def upconvolveReLU(in_channels, out_channels, kernel_size, stride, dim=3):
return nn.Sequential(ReLU, upconvolve(in_channels, out_channels, kernel_size, stride, dim=dim))
def upconvolveLeakyReLU(in_channels, out_channels, kernel_size, stride, dim=3):
return nn.Sequential(LeakyReLU(0.1), upconvolve(in_channels, out_channels, kernel_size, stride, dim=dim))
class VTN(nn.Module):
def init(self, dim=3, flow_multiplier=1., channels=16):
super(VTN, self).init()
self.flow_multiplier = flow_multiplier
self.channels = channels
self.dim = dim
self.conv1 = convolveLeakyReLU(2, channels, 3, 2, dim=dim)
self.conv2 = convolveLeakyReLU(channels, 2 * channels, 3, 2, dim=dim)
self.conv3 = convolveLeakyReLU(2 * channels, 4 * channels, 3, 2, dim=dim)
self.conv3_1 = convolveLeakyReLU(4 * channels, 4 * channels, 3, 1, dim=dim)
self.conv4 = convolveLeakyReLU(4 * channels, 8 * channels, 3, 2, dim=dim)
self.conv4_1 = convolveLeakyReLU(8 * channels, 8 * channels, 3, 1, dim=dim)
self.pred4 = convolve(8 * channels, dim, 3, 1, dim=dim)
self.upsamp4to3 = upconvolve(dim, dim, 4, 2, dim=dim)
self.deconv3 = upconvolveLeakyReLU(8 * channels, 4 * channels, 4, 2, dim=dim)
self.pred3 = convolve(8 * channels+dim, dim, 3, 1, dim=dim)
self.upsamp3to2 = upconvolve(dim, dim, 4, 2, dim=dim)
self.deconv2 = upconvolveLeakyReLU(8 * channels+dim, 2 * channels, 4, 2, dim=dim)
self.pred2 = convolve(4 * channels + dim, dim, 3, 1, dim=dim)
self.upsamp2to1 = upconvolve(dim, dim, 4, 2, dim=dim)
self.deconv1 = upconvolveLeakyReLU(4 * channels + dim, channels, 4, 2, dim=dim)
self.pred0 = upconvolve(2 * channels + dim, dim, 4, 2, dim=dim)
def forward(self, fixed, moving):
concat_image = torch.cat((fixed, moving), dim=1)
x1 = self.conv1(concat_image)
x2 = self.conv2(x1)
x3 = self.conv3(x2)
x3_1 = self.conv3_1(x3)
x4 = self.conv4(x3_1)
x4_1 = self.conv4_1(x4)
pred4 = self.pred4(x4_1)
upsamp4to3 = self.upsamp4to3(pred4)
deconv3 = self.deconv3(x4_1)
concat3 = torch.cat([x3_1, deconv3, upsamp4to3], dim=1)
pred3 = self.pred3(concat3)
upsamp3to2 = self.upsamp3to2(pred3)
deconv2 = self.deconv2(concat3)
concat2 = torch.cat([x2, deconv2, upsamp3to2], dim=1)
pred2 = self.pred2(concat2)
upsamp2to1 = self.upsamp2to1(pred2)
deconv1 = self.deconv1(concat2)
concat1 = torch.cat([x1, deconv1, upsamp2to1], dim=1)
pred0 = self.pred0(concat1)
return pred0 * 20 * self.flow_multiplier
class VTNAffineStem(nn.Module):
def __init__(self, dim=3, channels=16, flow_multiplier=1., im_size=128):
super(VTNAffineStem, self).__init__()
self.flow_multiplier = flow_multiplier
self.channels = channels
self.dim = dim
# Network architecture
# The first convolution's input is the concatenated image
self.conv1 = convolveLeakyReLU(2, channels, 3, 2, dim=self.dim)
self.conv2 = convolveLeakyReLU(channels, 2 * channels, 3, 2, dim=dim)
self.conv3 = convolveLeakyReLU(2 * channels, 4 * channels, 3, 2, dim=dim)
self.conv3_1 = convolveLeakyReLU(4 * channels, 4 * channels, 3, 1, dim=dim)
self.conv4 = convolveLeakyReLU(4 * channels, 8 * channels, 3, 2, dim=dim)
self.conv4_1 = convolveLeakyReLU(8 * channels, 8 * channels, 3, 1, dim=dim)
self.last_conv_size = im_size // (self.channels * 4)
#print(self.last_conv_size)
self.fc_loc = nn.Sequential(
nn.Linear(128 * self.last_conv_size**dim, 2048),
nn.ReLU(True),
nn.Dropout(0.5),
nn.Linear(2048, 1024),
nn.ReLU(True),
nn.Dropout(0.5),
nn.Linear(1024, 4),
nn.ReLU(True),
nn.Dropout(0.5),
nn.Linear(4, 6*(dim - 1))
)
self.fc_loc[-1].weight.data.zero_()
self.fc_loc[-1].bias.data.copy_(torch.tensor([1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0], dtype=torch.float))
def forward(self, fixed, moving):
#moving = torch.reshape(moving, fixed.shape)
concat_image = torch.cat((fixed, moving), dim=1) # 2 x 512 x 512
x1 = self.conv1(concat_image) # 16 x 256 x 256
x2 = self.conv2(x1) # 32 x 128 x 128
x3 = self.conv3(x2) # 1 x 64 x 64 x 64
x3_1 = self.conv3_1(x3) # 64 x 64 x 64
x4 = self.conv4(x3_1) # 128 x 32 x 32
x4_1 = self.conv4_1(x4) # 128 x 32 x 32
# Affine transformation
xs = x4_1.view(-1, 128 * self.last_conv_size ** self.dim)
theta = self.fc_loc(xs).view(-1, 3, 4)
flow = F.affine_grid(theta, moving.size(), align_corners=False) # batch x 512 x 512 x 2
flow = flow.permute(0, 4, 1, 2, 3)
return flow
if name == “main”:
vtn_model = VTN(dim=3)
vtn_affine_model = VTNAffineStem(dim=3, im_size=128)
x = torch.randn(1, 1, 90, 128, 128)
y1 = vtn_model(x, x)
y2 = vtn_affine_model(x, x)
assert y1.size() == y2.size()
