Upload ProbUNet

PULASki.py

@@ -0,0 +1,48 @@
+
+import sys
+
+import torch
+import torch.nn as nn
+
+from transformers import PreTrainedModel
+
+from .ProbUNet_model import InjectionConvEncoder2D, InjectionUNet2D, InjectionConvEncoder3D, InjectionUNet3D, ProbabilisticSegmentationNet
+from .PULASkiConfigs import ProbUNetConfig
+
+class ProbUNet(PreTrainedModel):
+    config_class = ProbUNetConfig
+    def __init__(self, config):
+        super().__init__(config)
+        
+        if config.dim == 2:
+            task_op = InjectionUNet2D
+            prior_op = InjectionConvEncoder2D
+            posterior_op = InjectionConvEncoder2D
+        elif config.dim == 3:
+            task_op = InjectionUNet3D
+            prior_op = InjectionConvEncoder3D
+            posterior_op = InjectionConvEncoder3D
+        else:
+            sys.exit("Invalid dim! Only configured for dim 2 and 3.")
+            
+        if config.latent_distribution == "normal":
+            latent_distribution = torch.distributions.Normal
+        else:
+            sys.exit("Invalid latent_distribution. Only normal has been implemented.")
+            
+        self.model = ProbabilisticSegmentationNet(in_channels=config.in_channels, 
+                                                    out_channels=config.out_channels, 
+                                                    num_feature_maps=config.num_feature_maps,
+                                                    latent_size=config.latent_size,
+                                                    depth=config.depth,
+                                                    latent_distribution=latent_distribution,
+                                                    task_op=task_op,
+                                                    task_kwargs={"output_activation_op": nn.Identity if config.no_outact_op else nn.Sigmoid, 
+                                                                    "activation_kwargs": {"inplace": True}, "injection_at":  config.prob_injection_at},  
+                                                    prior_op=prior_op,
+                                                    prior_kwargs={"activation_kwargs": {"inplace": True}, "norm_depth": 2}, 
+                                                    posterior_op=posterior_op,
+                                                    posterior_kwargs={"activation_kwargs": {"inplace": True}, "norm_depth": 2},
+                                                )
+    def forward(self, x):
+        return self.model(x)

PULASkiConfigs.py

@@ -0,0 +1,26 @@
+from transformers import PretrainedConfig
+
+class ProbUNetConfig(PretrainedConfig):
+    model_type = "ProbUNet"
+    def __init__(
+            self,
+            dim=2,
+            in_channels=1,
+            out_channels=1,
+            num_feature_maps=24,
+            latent_size=3,
+            depth=5,
+            latent_distribution="normal",
+            no_outact_op=False,
+            prob_injection_at="end",
+            **kwargs):
+        self.dim = dim
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.num_feature_maps = num_feature_maps
+        self.latent_size = latent_size
+        self.depth = depth
+        self.latent_distribution = latent_distribution
+        self.no_outact_op = no_outact_op
+        self.prob_injection_at = prob_injection_at
+        super().__init__(**kwargs)

ProbUNet_model.py

@@ -0,0 +1,731 @@
+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)

ProbUNet_utils.py

@@ -0,0 +1,224 @@
+import argparse
+import os
+import shutil
+import numpy as np
+import torch
+# from trixi.util import Config, GridSearch
+
+
+def check_attributes(object_, attributes):
+
+    missing = []
+    for attr in attributes:
+        if not hasattr(object_, attr):
+            missing.append(attr)
+    if len(missing) > 0:
+        return False
+    else:
+        return True
+
+
+def set_seeds(seed, cuda=True):
+    if not hasattr(seed, "__iter__"):
+        seed = (seed, seed, seed)
+    np.random.seed(seed[0])
+    torch.manual_seed(seed[1])
+    if cuda: torch.cuda.manual_seed_all(seed[2])
+
+
+def make_onehot(array, labels=None, axis=1, newaxis=False):
+
+    # get labels if necessary
+    if labels is None:
+        labels = np.unique(array)
+        labels = list(map(lambda x: x.item(), labels))
+
+    # get target shape
+    new_shape = list(array.shape)
+    if newaxis:
+        new_shape.insert(axis, len(labels))
+    else:
+        new_shape[axis] = new_shape[axis] * len(labels)
+
+    # make zero array
+    if type(array) == np.ndarray:
+        new_array = np.zeros(new_shape, dtype=array.dtype)
+    elif torch.is_tensor(array):
+        new_array = torch.zeros(new_shape, dtype=array.dtype, device=array.device)
+    else:
+        raise TypeError("Onehot conversion undefined for object of type {}".format(type(array)))
+
+    # fill new array
+    n_seg_channels = 1 if newaxis else array.shape[axis]
+    for seg_channel in range(n_seg_channels):
+        for l, label in enumerate(labels):
+            new_slc = [slice(None), ] * len(new_shape)
+            slc = [slice(None), ] * len(array.shape)
+            new_slc[axis] = seg_channel * len(labels) + l
+            if not newaxis:
+                slc[axis] = seg_channel
+            new_array[tuple(new_slc)] = array[tuple(slc)] == label
+
+    return new_array
+
+
+def match_to(x, ref, keep_axes=(1,)):
+
+    target_shape = list(ref.shape)
+    for i in keep_axes:
+        target_shape[i] = x.shape[i]
+    target_shape = tuple(target_shape)
+    if x.shape == target_shape:
+        pass
+    if x.dim() == 1:
+        x = x.unsqueeze(0)
+    if x.dim() == 2:
+        while x.dim() < len(target_shape):
+            x = x.unsqueeze(-1)
+
+    x = x.expand(*target_shape)
+    x = x.to(device=ref.device, dtype=ref.dtype)
+
+    return x
+
+
+def make_slices(original_shape, patch_shape):
+
+    working_shape = original_shape[-len(patch_shape):]
+    splits = []
+    for i in range(len(working_shape)):
+        splits.append([])
+        for j in range(working_shape[i] // patch_shape[i]):
+            splits[i].append(slice(j*patch_shape[i], (j+1)*patch_shape[i]))
+        rest = working_shape[i] % patch_shape[i]
+        if rest > 0:
+            splits[i].append(slice((j+1)*patch_shape[i], (j+1)*patch_shape[i] + rest))
+
+    # now we have all slices for the individual dimensions
+    # we need their combinatorial combinations
+    slices = list(itertools.product(*splits))
+    for i in range(len(slices)):
+        slices[i] = [slice(None), ] * (len(original_shape) - len(patch_shape)) + list(slices[i])
+
+    return slices
+
+
+def coordinate_grid_samples(mean, std, factor_std=5, scale_std=1.):
+
+    relative = np.linspace(-scale_std*factor_std, scale_std*factor_std, 2*factor_std+1)
+    positions = np.array([mean + i * std for i in relative]).T
+    axes = np.meshgrid(*positions)
+    axes = map(lambda x: list(x.ravel()), axes)
+    samples = list(zip(*axes))
+    samples = list(map(np.array, samples))
+
+    return samples
+
+
+def get_default_experiment_parser():
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("base_dir", type=str, help="Working directory for experiment.")
+    parser.add_argument("-c", "--config", type=str, default=None, help="Path to a config file.")
+    parser.add_argument("-v", "--visdomlogger", action="store_true", help="Use visdomlogger.")
+    parser.add_argument("-tx", "--tensorboardxlogger", type=str, default=None)
+    parser.add_argument("-tl", "--telegramlogger", action="store_true")
+    parser.add_argument("-dc", "--default_config", type=str, default="DEFAULTS", help="Select a default Config")
+    parser.add_argument("-ad", "--automatic_description", action="store_true")
+    parser.add_argument("-r", "--resume", type=str, default=None, help="Path to resume from")
+    parser.add_argument("-irc", "--ignore_resume_config", action="store_true", help="Ignore Config in experiment we resume from.")
+    parser.add_argument("-test", "--test", action="store_true", help="Run test instead of training")
+    parser.add_argument("-g", "--grid", type=str, help="Path to a config for grid search")
+    parser.add_argument("-s", "--skip_existing", action="store_true", help="Skip configs for which an experiment exists, only for grid search")
+    parser.add_argument("-m", "--mods", type=str, nargs="+", default=None, help="Mods are Config stubs to update only relevant parts for a certain setup.")
+    parser.add_argument("-ct", "--copy_test", action="store_true", help="Copy test files to original experiment.")
+
+    return parser
+
+
+def run_experiment(experiment, configs, args, mods=None, **kwargs):
+
+    # set a few defaults
+    if "explogger_kwargs" not in kwargs:
+        kwargs["explogger_kwargs"] = dict(folder_format="{experiment_name}_%Y%m%d-%H%M%S")
+    if "explogger_freq" not in kwargs:
+        kwargs["explogger_freq"] = 1
+    if "resume_save_types" not in kwargs:
+        kwargs["resume_save_types"] = ("model", "simple", "th_vars", "results")
+
+    config = Config(file_=args.config) if args.config is not None else Config()
+    config.update_missing(configs[args.default_config].deepcopy())
+    if args.mods is not None and mods is not None:
+        for mod in args.mods:
+            config.update(mods[mod])
+    config = Config(config=config, update_from_argv=True)
+
+    # GET EXISTING EXPERIMENTS TO BE ABLE TO SKIP CERTAIN CONFIGS
+    if args.skip_existing:
+        existing_configs = []
+        for exp in os.listdir(args.base_dir):
+            try:
+                existing_configs.append(Config(file_=os.path.join(args.base_dir, exp, "config", "config.json")))
+            except Exception as e:
+                pass
+
+    if args.grid is not None:
+        grid = GridSearch().read(args.grid)
+    else:
+        grid = [{}]
+
+    for combi in grid:
+
+        config.update(combi)
+
+        if args.skip_existing:
+            skip_this = False
+            for existing_config in existing_configs:
+                if existing_config.contains(config):
+                    skip_this = True
+                    break
+            if skip_this:
+                continue
+
+        if "backup_every" in config:
+            kwargs["save_checkpoint_every_epoch"] = config["backup_every"]
+
+        loggers = {}
+        if args.visdomlogger:
+            loggers["v"] = ("visdom", {}, 1)
+        if args.tensorboardxlogger is not None:
+            if args.tensorboardxlogger == "same":
+                loggers["tx"] = ("tensorboard", {}, 1)
+            else:
+                loggers["tx"] = ("tensorboard", {"target_dir": args.tensorboardxlogger}, 1)
+
+        if args.telegramlogger:
+            kwargs["use_telegram"] = True
+
+        if args.automatic_description:
+            difference_to_default = Config.difference_config_static(config, configs["DEFAULTS"]).flat(keep_lists=True, max_split_size=0, flatten_int=True)
+            description_str = ""
+            for key, val in difference_to_default.items():
+                val = val[0]
+                description_str = "{} = {}\n{}".format(key, val, description_str)
+            config.description = description_str
+
+        exp = experiment(config=config,
+                         base_dir=args.base_dir,
+                         resume=args.resume,
+                         ignore_resume_config=args.ignore_resume_config,
+                         loggers=loggers,
+                         **kwargs)
+
+        trained = False
+        if args.resume is None or args.test is False:
+            exp.run()
+            trained = True
+        if args.test:
+            exp.run_test(setup=not trained)
+            if isinstance(args.resume, str) and exp.elog is not None and args.copy_test:
+                for f in glob.glob(os.path.join(exp.elog.save_dir, "test*")):
+                    if os.path.isdir(f):
+                        shutil.copytree(f, os.path.join(args.resume, "save", os.path.basename(f)))
+                    else:
+                        shutil.copy(f, os.path.join(args.resume, "save"))

config.json

@@ -0,0 +1,21 @@
+{
+  "architectures": [
+    "ProbUNet"
+  ],
+  "auto_map": {
+    "AutoConfig": "PULASkiConfigs.ProbUNetConfig",
+    "AutoModel": "PULASki.ProbUNet"
+  },
+  "depth": 5,
+  "dim": 2,
+  "in_channels": 1,
+  "latent_distribution": "normal",
+  "latent_size": 3,
+  "model_type": "ProbUNet",
+  "no_outact_op": false,
+  "num_feature_maps": 24,
+  "out_channels": 1,
+  "prob_injection_at": "end",
+  "torch_dtype": "float32",
+  "transformers_version": "4.44.2"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f14fa2e9b8ba02d99a5eaf22a2d72af7ef1cd63e42aa0b136d23b75033518fc
+size 40660760