import torch
import torch.nn as nn
from .ProbUNet_utils import make_onehot as make_onehot_segmentation, make_slices, match_to
def is_conv(op):
conv_types = (nn.Conv1d,
nn.Conv2d,
nn.Conv3d,
nn.ConvTranspose1d,
nn.ConvTranspose2d,
nn.ConvTranspose3d)
if type(op) == type and issubclass(op, conv_types):
return True
elif type(op) in conv_types:
return True
else:
return False
class ConvModule(nn.Module):
def __init__(self, *args, **kwargs):
super(ConvModule, self).__init__()
def init_weights(self, init_fn, *args, **kwargs):
class init_(object):
def __init__(self):
self.fn = init_fn
self.args = args
self.kwargs = kwargs
def __call__(self, module):
if is_conv(type(module)):
module.weight = self.fn(module.weight, *self.args, **self.kwargs)
_init_ = init_()
self.apply(_init_)
def init_bias(self, init_fn, *args, **kwargs):
class init_(object):
def __init__(self):
self.fn = init_fn
self.args = args
self.kwargs = kwargs
def __call__(self, module):
if is_conv(type(module)) and module.bias is not None:
module.bias = self.fn(module.bias, *self.args, **self.kwargs)
_init_ = init_()
self.apply(_init_)
class ConcatCoords(nn.Module):
def forward(self, input_):
dim = input_.dim() - 2
coord_channels = []
for i in range(dim):
view = [1, ] * dim
view[i] = -1
repeat = list(input_.shape[2:])
repeat[i] = 1
coord_channels.append(
torch.linspace(-0.5, 0.5, input_.shape[i+2])
.view(*view)
.repeat(*repeat)
.to(device=input_.device, dtype=input_.dtype))
coord_channels = torch.stack(coord_channels).unsqueeze(0)
repeat = [1, ] * input_.dim()
repeat[0] = input_.shape[0]
coord_channels = coord_channels.repeat(*repeat).contiguous()
return torch.cat([input_, coord_channels], 1)
class InjectionConvEncoder(ConvModule):
_default_activation_kwargs = dict(inplace=True)
_default_norm_kwargs = dict()
_default_conv_kwargs = dict(kernel_size=3, padding=1)
_default_pool_kwargs = dict(kernel_size=2)
_default_dropout_kwargs = dict()
_default_global_pool_kwargs = dict()
def __init__(self,
in_channels=1,
out_channels=6,
depth=4,
injection_depth="last",
injection_channels=0,
block_depth=2,
num_feature_maps=24,
feature_map_multiplier=2,
activation_op=nn.LeakyReLU,
activation_kwargs=None,
norm_op=nn.InstanceNorm2d,
norm_kwargs=None,
norm_depth=0,
conv_op=nn.Conv2d,
conv_kwargs=None,
pool_op=nn.AvgPool2d,
pool_kwargs=None,
dropout_op=None,
dropout_kwargs=None,
global_pool_op=nn.AdaptiveAvgPool2d,
global_pool_kwargs=None,
**kwargs):
super(InjectionConvEncoder, self).__init__(**kwargs)
self.in_channels = in_channels
self.out_channels = out_channels
self.depth = depth
self.injection_depth = depth - 1 if injection_depth == "last" else injection_depth
self.injection_channels = injection_channels
self.block_depth = block_depth
self.num_feature_maps = num_feature_maps
self.feature_map_multiplier = feature_map_multiplier
self.activation_op = activation_op
self.activation_kwargs = self._default_activation_kwargs
if activation_kwargs is not None:
self.activation_kwargs.update(activation_kwargs)
self.norm_op = norm_op
self.norm_kwargs = self._default_norm_kwargs
if norm_kwargs is not None:
self.norm_kwargs.update(norm_kwargs)
self.norm_depth = depth if norm_depth == "full" else norm_depth
self.conv_op = conv_op
self.conv_kwargs = self._default_conv_kwargs
if conv_kwargs is not None:
self.conv_kwargs.update(conv_kwargs)
self.pool_op = pool_op
self.pool_kwargs = self._default_pool_kwargs
if pool_kwargs is not None:
self.pool_kwargs.update(pool_kwargs)
self.dropout_op = dropout_op
self.dropout_kwargs = self._default_dropout_kwargs
if dropout_kwargs is not None:
self.dropout_kwargs.update(dropout_kwargs)
self.global_pool_op = global_pool_op
self.global_pool_kwargs = self._default_global_pool_kwargs
if global_pool_kwargs is not None:
self.global_pool_kwargs.update(global_pool_kwargs)
for d in range(self.depth):
in_ = self.in_channels if d == 0 else self.num_feature_maps * (self.feature_map_multiplier**(d-1))
out_ = self.num_feature_maps * (self.feature_map_multiplier**d)
if d == self.injection_depth + 1:
in_ += self.injection_channels
layers = []
if d > 0:
layers.append(self.pool_op(**self.pool_kwargs))
for b in range(self.block_depth):
current_in = in_ if b == 0 else out_
layers.append(self.conv_op(current_in, out_, **self.conv_kwargs))
if self.norm_op is not None and d < self.norm_depth:
layers.append(self.norm_op(out_, **self.norm_kwargs))
if self.activation_op is not None:
layers.append(self.activation_op(**self.activation_kwargs))
if self.dropout_op is not None:
layers.append(self.dropout_op(**self.dropout_kwargs))
if d == self.depth - 1:
current_conv_kwargs = self.conv_kwargs.copy()
current_conv_kwargs["kernel_size"] = 1
current_conv_kwargs["padding"] = 0
current_conv_kwargs["bias"] = False
layers.append(self.conv_op(out_, out_channels, **current_conv_kwargs))
self.add_module("encode_{}".format(d), nn.Sequential(*layers))
if self.global_pool_op is not None:
self.add_module("global_pool", self.global_pool_op(1, **self.global_pool_kwargs))
def forward(self, x, injection=None):
for d in range(self.depth):
x = self._modules["encode_{}".format(d)](x)
if d == self.injection_depth and self.injection_channels > 0:
injection = match_to(injection, x, self.injection_channels)
x = torch.cat([x, injection], 1)
if hasattr(self, "global_pool"):
x = self.global_pool(x)
return x
class InjectionConvEncoder3D(InjectionConvEncoder):
def __init__(self, *args, **kwargs):
update_kwargs = dict(
norm_op=nn.InstanceNorm3d,
conv_op=nn.Conv3d,
pool_op=nn.AvgPool3d,
global_pool_op=nn.AdaptiveAvgPool3d
)
for (arg, val) in update_kwargs.items():
if arg not in kwargs: kwargs[arg] = val
super(InjectionConvEncoder3D, self).__init__(*args, **kwargs)
class InjectionConvEncoder2D(InjectionConvEncoder): #Created by Soumick
def __init__(self, *args, **kwargs):
update_kwargs = dict(
norm_op=nn.InstanceNorm2d,
conv_op=nn.Conv2d,
pool_op=nn.AvgPool2d,
global_pool_op=nn.AdaptiveAvgPool2d
)
for (arg, val) in update_kwargs.items():
if arg not in kwargs: kwargs[arg] = val
super(InjectionConvEncoder2D, self).__init__(*args, **kwargs)
class InjectionUNet(ConvModule):
def __init__(
self,
depth=5,
in_channels=4,
out_channels=4,
kernel_size=3,
dilation=1,
num_feature_maps=24,
block_depth=2,
num_1x1_at_end=3,
injection_channels=3,
injection_at="end",
activation_op=nn.LeakyReLU,
activation_kwargs=None,
pool_op=nn.AvgPool2d,
pool_kwargs=dict(kernel_size=2),
dropout_op=None,
dropout_kwargs=None,
norm_op=nn.InstanceNorm2d,
norm_kwargs=None,
conv_op=nn.Conv2d,
conv_kwargs=None,
upconv_op=nn.ConvTranspose2d,
upconv_kwargs=None,
output_activation_op=None,
output_activation_kwargs=None,
return_bottom=False,
coords=False,
coords_dim=2,
**kwargs
):
super(InjectionUNet, self).__init__(**kwargs)
self.depth = depth
self.in_channels = in_channels
self.out_channels = out_channels
self.kernel_size = kernel_size
self.dilation = dilation
self.padding = (self.kernel_size + (self.kernel_size-1) * (self.dilation-1)) // 2
self.num_feature_maps = num_feature_maps
self.block_depth = block_depth
self.num_1x1_at_end = num_1x1_at_end
self.injection_channels = injection_channels
self.injection_at = injection_at
self.activation_op = activation_op
self.activation_kwargs = {} if activation_kwargs is None else activation_kwargs
self.pool_op = pool_op
self.pool_kwargs = {} if pool_kwargs is None else pool_kwargs
self.dropout_op = dropout_op
self.dropout_kwargs = {} if dropout_kwargs is None else dropout_kwargs
self.norm_op = norm_op
self.norm_kwargs = {} if norm_kwargs is None else norm_kwargs
self.conv_op = conv_op
self.conv_kwargs = {} if conv_kwargs is None else conv_kwargs
self.upconv_op = upconv_op
self.upconv_kwargs = {} if upconv_kwargs is None else upconv_kwargs
self.output_activation_op = output_activation_op
self.output_activation_kwargs = {} if output_activation_kwargs is None else output_activation_kwargs
self.return_bottom = return_bottom
if not coords:
self.coords = [[], []]
elif coords is True:
self.coords = [list(range(depth)), []]
else:
self.coords = coords
self.coords_dim = coords_dim
self.last_activations = None
# BUILD ENCODER
for d in range(self.depth):
block = []
if d > 0:
block.append(self.pool_op(**self.pool_kwargs))
for i in range(self.block_depth):
# bottom block fixed to have depth 1
if d == self.depth - 1 and i > 0:
continue
out_size = self.num_feature_maps * 2**d
if d == 0 and i == 0:
in_size = self.in_channels
elif i == 0:
in_size = self.num_feature_maps * 2**(d - 1)
else:
in_size = out_size
# check for coord appending at this depth
if d in self.coords[0] and i == 0:
block.append(ConcatCoords())
in_size += self.coords_dim
block.append(self.conv_op(in_size,
out_size,
self.kernel_size,
padding=self.padding,
dilation=self.dilation,
**self.conv_kwargs))
if self.dropout_op is not None:
block.append(self.dropout_op(**self.dropout_kwargs))
if self.norm_op is not None:
block.append(self.norm_op(out_size, **self.norm_kwargs))
block.append(self.activation_op(**self.activation_kwargs))
self.add_module("encode-{}".format(d), nn.Sequential(*block))
# BUILD DECODER
for d in reversed(range(self.depth)):
block = []
for i in range(self.block_depth):
# bottom block fixed to have depth 1
if d == self.depth - 1 and i > 0:
continue
out_size = self.num_feature_maps * 2**(d)
if i == 0 and d < self.depth - 1:
in_size = self.num_feature_maps * 2**(d+1)
elif i == 0 and self.injection_at == "bottom":
in_size = out_size + self.injection_channels
else:
in_size = out_size
# check for coord appending at this depth
if d in self.coords[0] and i == 0 and d < self.depth - 1:
block.append(ConcatCoords())
in_size += self.coords_dim
block.append(self.conv_op(in_size,
out_size,
self.kernel_size,
padding=self.padding,
dilation=self.dilation,
**self.conv_kwargs))
if self.dropout_op is not None:
block.append(self.dropout_op(**self.dropout_kwargs))
if self.norm_op is not None:
block.append(self.norm_op(out_size, **self.norm_kwargs))
block.append(self.activation_op(**self.activation_kwargs))
if d > 0:
block.append(self.upconv_op(out_size,
out_size // 2,
self.kernel_size,
2,
padding=self.padding,
dilation=self.dilation,
output_padding=1,
**self.upconv_kwargs))
self.add_module("decode-{}".format(d), nn.Sequential(*block))
if self.injection_at == "end":
out_size += self.injection_channels
in_size = out_size
for i in range(self.num_1x1_at_end):
if i == self.num_1x1_at_end - 1:
out_size = self.out_channels
current_conv_kwargs = self.conv_kwargs.copy()
current_conv_kwargs["bias"] = True
self.add_module("reduce-{}".format(i), self.conv_op(in_size, out_size, 1, **current_conv_kwargs))
if i != self.num_1x1_at_end - 1:
self.add_module("reduce-{}-nonlin".format(i), self.activation_op(**self.activation_kwargs))
if self.output_activation_op is not None:
self.add_module("output-activation", self.output_activation_op(**self.output_activation_kwargs))
def reset(self):
self.last_activations = None
def forward(self, x, injection=None, reuse_last_activations=False, store_activations=False):
if self.injection_at == "bottom": # not worth it for now
reuse_last_activations = False
store_activations = False
if self.last_activations is None or reuse_last_activations is False:
enc = [x]
for i in range(self.depth - 1):
enc.append(self._modules["encode-{}".format(i)](enc[-1]))
bottom_rep = self._modules["encode-{}".format(self.depth - 1)](enc[-1])
if self.injection_at == "bottom" and self.injection_channels > 0:
injection = match_to(injection, bottom_rep, (0, 1))
bottom_rep = torch.cat((bottom_rep, injection), 1)
x = self._modules["decode-{}".format(self.depth - 1)](bottom_rep)
for i in reversed(range(self.depth - 1)):
x = self._modules["decode-{}".format(i)](torch.cat((enc[-(self.depth - 1 - i)], x), 1))
if store_activations:
self.last_activations = x.detach()
else:
x = self.last_activations
if self.injection_at == "end" and self.injection_channels > 0:
injection = match_to(injection, x, (0, 1))
x = torch.cat((x, injection), 1)
for i in range(self.num_1x1_at_end):
x = self._modules["reduce-{}".format(i)](x)
if self.output_activation_op is not None:
x = self._modules["output-activation"](x)
if self.return_bottom and not reuse_last_activations:
return x, bottom_rep
else:
return x
class InjectionUNet3D(InjectionUNet):
def __init__(self, *args, **kwargs):
update_kwargs = dict(
pool_op=nn.AvgPool3d,
norm_op=nn.InstanceNorm3d,
conv_op=nn.Conv3d,
upconv_op=nn.ConvTranspose3d,
coords_dim=3
)
for (arg, val) in update_kwargs.items():
if arg not in kwargs: kwargs[arg] = val
super(InjectionUNet3D, self).__init__(*args, **kwargs)
class InjectionUNet2D(InjectionUNet): #Created by Soumick
def __init__(self, *args, **kwargs):
update_kwargs = dict(
pool_op=nn.AvgPool2d,
norm_op=nn.InstanceNorm2d,
conv_op=nn.Conv2d,
upconv_op=nn.ConvTranspose2d,
coords_dim=2
)
for (arg, val) in update_kwargs.items():
if arg not in kwargs: kwargs[arg] = val
super(InjectionUNet2D, self).__init__(*args, **kwargs)
class ProbabilisticSegmentationNet(ConvModule):
def __init__(self,
in_channels=4,
out_channels=4,
num_feature_maps=24,
latent_size=3,
depth=5,
latent_distribution=torch.distributions.Normal,
task_op=InjectionUNet3D,
task_kwargs=None,
prior_op=InjectionConvEncoder3D,
prior_kwargs=None,
posterior_op=InjectionConvEncoder3D,
posterior_kwargs=None,
**kwargs):
super(ProbabilisticSegmentationNet, self).__init__(**kwargs)
self.task_op = task_op
self.task_kwargs = {} if task_kwargs is None else task_kwargs
self.prior_op = prior_op
self.prior_kwargs = {} if prior_kwargs is None else prior_kwargs
self.posterior_op = posterior_op
self.posterior_kwargs = {} if posterior_kwargs is None else posterior_kwargs
default_task_kwargs = dict(
in_channels=in_channels,
out_channels=out_channels,
num_feature_maps=num_feature_maps,
injection_size=latent_size,
depth=depth
)
default_prior_kwargs = dict(
in_channels=in_channels,
out_channels=latent_size*2, #Soumick
num_feature_maps=num_feature_maps,
z_dim=latent_size,
depth=depth
)
default_posterior_kwargs = dict(
in_channels=in_channels+out_channels,
out_channels=latent_size*2, #Soumick
num_feature_maps=num_feature_maps,
z_dim=latent_size,
depth=depth
)
default_task_kwargs.update(self.task_kwargs)
self.task_kwargs = default_task_kwargs
default_prior_kwargs.update(self.prior_kwargs)
self.prior_kwargs = default_prior_kwargs
default_posterior_kwargs.update(self.posterior_kwargs)
self.posterior_kwargs = default_posterior_kwargs
self.latent_distribution = latent_distribution
self._prior = None
self._posterior = None
self.make_modules()
def make_modules(self):
if type(self.task_op) == type:
self.add_module("task_net", self.task_op(**self.task_kwargs))
else:
self.add_module("task_net", self.task_op)
if type(self.prior_op) == type:
self.add_module("prior_net", self.prior_op(**self.prior_kwargs))
else:
self.add_module("prior_net", self.prior_op)
if type(self.posterior_op) == type:
self.add_module("posterior_net", self.posterior_op(**self.posterior_kwargs))
else:
self.add_module("posterior_net", self.posterior_op)
@property
def prior(self):
return self._prior
@property
def posterior(self):
return self._posterior
@property
def last_activations(self):
return self.task_net.last_activations
def train(self, mode=True):
super(ProbabilisticSegmentationNet, self).train(mode)
self.reset()
def reset(self):
self.task_net.reset()
self._prior = None
self._posterior = None
def forward(self, input_, seg=None, make_onehot=True, make_onehot_classes=None, newaxis=False, distlossN=0):
"""Forward pass includes reparametrization sampling during training, otherwise it'll just take the prior mean."""
self.encode_prior(input_)
if distlossN == 0:
if self.training:
self.encode_posterior(input_, seg, make_onehot, make_onehot_classes, newaxis)
sample = self.posterior.rsample()
else:
sample = self.prior.loc
return self.task_net(input_, sample, store_activations=not self.training)
else:
if self.training:
self.encode_posterior(input_, seg, make_onehot, make_onehot_classes, newaxis)
segs = []
for i in range(distlossN):
sample = self.posterior.rsample()
segs.append(self.task_net(input_, sample, store_activations=not self.training))
return segs #torch.concat(segs, dim=0)
else: #I'm not totally sure about this!!
sample = self.prior.loc
return self.task_net(input_, sample, store_activations=not self.training)
def encode_prior(self, input_):
rep = self.prior_net(input_)
if isinstance(rep, tuple):
mean, logvar = rep
elif torch.is_tensor(rep):
mean, logvar = torch.split(rep, rep.shape[1] // 2, dim=1)
self._prior = self.latent_distribution(mean, logvar.mul(0.5).exp())
return self._prior
def encode_posterior(self, input_, seg, make_onehot=True, make_onehot_classes=None, newaxis=False):
if make_onehot:
if make_onehot_classes is None:
make_onehot_classes = tuple(range(self.posterior_net.in_channels - input_.shape[1]))
seg = make_onehot_segmentation(seg, make_onehot_classes, newaxis=newaxis)
rep = self.posterior_net(torch.cat((input_, seg.float()), 1))
if isinstance(rep, tuple):
mean, logvar = rep
elif torch.is_tensor(rep):
mean, logvar = torch.split(rep, rep.shape[1] // 2, dim=1)
self._posterior = self.latent_distribution(mean, logvar.mul(0.5).exp())
return self._posterior
def sample_prior(self, N=1, out_device=None, input_=None, pred_with_mean=False):
"""Draw multiple samples from the current prior.
* input_ is required if no activations are stored in task_net.
* If input_ is given, prior will automatically be encoded again.
* Returns either a single sample or a list of samples.
"""
if out_device is None:
if self.last_activations is not None:
out_device = self.last_activations.device
elif input_ is not None:
out_device = input_.device
else:
out_device = next(self.task_net.parameters()).device
with torch.no_grad():
if self.prior is None or input_ is not None:
self.encode_prior(input_)
result = []
if input_ is not None:
result.append(self.task_net(input_, self.prior.sample(), reuse_last_activations=False, store_activations=True).to(device=out_device))
while len(result) < N:
result.append(self.task_net(input_,
self.prior.sample(),
reuse_last_activations=self.last_activations is not None,
store_activations=False).to(device=out_device))
if pred_with_mean:
result.append(self.task_net(input_, self.prior.mean, reuse_last_activations=False, store_activations=True).to(device=out_device))
if len(result) == 1:
return result[0]
else:
return result
def reconstruct(self, sample=None, use_posterior_mean=True, out_device=None, input_=None):
"""Reconstruct a sample or the current posterior mean. Will not compute gradients!"""
if self.posterior is None and sample is None:
raise ValueError("'posterior' is currently None. Please pass an input and a segmentation first.")
if out_device is None:
out_device = next(self.task_net.parameters()).device
if sample is None:
if use_posterior_mean:
sample = self.posterior.loc
else:
sample = self.posterior.sample()
else:
sample = sample.to(next(self.task_net.parameters()).device)
with torch.no_grad():
return self.task_net(input_, sample, reuse_last_activations=True).to(device=out_device)
def kl_divergence(self):
"""Compute current KL, requires existing prior and posterior."""
if self.posterior is None or self.prior is None:
raise ValueError("'prior' and 'posterior' must not be None, but prior={} and posterior={}".format(self.prior, self.posterior))
return torch.distributions.kl_divergence(self.posterior, self.prior).sum()
def elbo(self, seg, input_=None, nll_reduction="sum", beta=1.0, make_onehot=True, make_onehot_classes=None, newaxis=False):
"""Compute the ELBO with seg as ground truth.
* Prior is expected and will not be encoded.
* If input_ is given, posterior will automatically be encoded.
* Either input_ or stored activations must be available.
"""
if self.last_activations is None:
raise ValueError("'last_activations' is currently None. Please pass an input first.")
if input_ is not None:
with torch.no_grad():
self.encode_posterior(input_, seg, make_onehot=make_onehot, make_onehot_classes=make_onehot_classes, newaxis=newaxis)
if make_onehot and newaxis:
pass # seg will already be (B x SPACE)
elif make_onehot and not newaxis:
seg = seg[:, 0] # in this case seg will hopefully be (B x 1 x SPACE)
else:
seg = torch.argmax(seg, 1, keepdim=False) # seg is already onehot
kl = self.kl_divergence()
nll = nn.NLLLoss(reduction=nll_reduction)(self.reconstruct(sample=None, use_posterior_mean=True, out_device=None), seg.long())
return - (beta * nll + kl)