File size: 27,509 Bytes
28b6cf8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
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)