Unet.py

#    Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.

from copy import deepcopy
from torch import nn
import torch.nn.functional as F
import torch
import numpy as np
import torch.nn.functional

def softmax_helper(x):
    rpt = [1 for _ in range(len(x.size()))]
    rpt[1] = x.size(1)
    x_max = x.max(1, keepdim=True)[0].repeat(*rpt)
    e_x = torch.exp(x - x_max)
    return e_x / e_x.sum(1, keepdim=True).repeat(*rpt)

class InitWeights_He(object):
    def __init__(self, neg_slope=1e-2):
        self.neg_slope = neg_slope

    def __call__(self, module):
        if isinstance(module, nn.Conv3d) or isinstance(module, nn.Conv2d) or isinstance(module, nn.ConvTranspose2d) or isinstance(module, nn.ConvTranspose3d):
            module.weight = nn.init.kaiming_normal_(module.weight, a=self.neg_slope)
            if module.bias is not None:
                module.bias = nn.init.constant_(module.bias, 0)

class ConvDropoutNormNonlin(nn.Module):
    """
    fixes a bug in ConvDropoutNormNonlin where lrelu was used regardless of nonlin. Bad.
    """

    def __init__(self, input_channels, output_channels,
                 conv_op=nn.Conv2d, conv_kwargs=None,
                 norm_op=nn.BatchNorm2d, norm_op_kwargs=None,
                 dropout_op=nn.Dropout2d, dropout_op_kwargs=None,
                 nonlin=nn.LeakyReLU, nonlin_kwargs=None):
        super(ConvDropoutNormNonlin, self).__init__()
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}
        if conv_kwargs is None:
            conv_kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1, 'dilation': 1, 'bias': True}

        self.nonlin_kwargs = nonlin_kwargs
        self.nonlin = nonlin
        self.dropout_op = dropout_op
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.conv_kwargs = conv_kwargs
        self.conv_op = conv_op
        self.norm_op = norm_op

        self.conv = self.conv_op(input_channels, output_channels, **self.conv_kwargs)
        if self.dropout_op is not None and self.dropout_op_kwargs['p'] is not None and self.dropout_op_kwargs[
            'p'] > 0:
            self.dropout = self.dropout_op(**self.dropout_op_kwargs)
        else:
            self.dropout = None
        self.instnorm = self.norm_op(output_channels, **self.norm_op_kwargs)
        self.lrelu = self.nonlin(**self.nonlin_kwargs)

    def forward(self, x):
        x = self.conv(x)
        if self.dropout is not None:
            x = self.dropout(x)
        return self.lrelu(self.instnorm(x))


class ConvDropoutNonlinNorm(ConvDropoutNormNonlin):
    def forward(self, x):
        x = self.conv(x)
        if self.dropout is not None:
            x = self.dropout(x)
        return self.instnorm(self.lrelu(x))


class StackedConvLayers(nn.Module):
    def __init__(self, input_feature_channels, output_feature_channels, num_convs,
                 conv_op=nn.Conv2d, conv_kwargs=None,
                 norm_op=nn.BatchNorm2d, norm_op_kwargs=None,
                 dropout_op=nn.Dropout2d, dropout_op_kwargs=None,
                 nonlin=nn.LeakyReLU, nonlin_kwargs=None, first_stride=None, basic_block=ConvDropoutNormNonlin):
        '''
        stacks ConvDropoutNormLReLU layers. initial_stride will only be applied to first layer in the stack. The other parameters affect all layers
        :param input_feature_channels:
        :param output_feature_channels:
        :param num_convs:
        :param dilation:
        :param kernel_size:
        :param padding:
        :param dropout:
        :param initial_stride:
        :param conv_op:
        :param norm_op:
        :param dropout_op:
        :param inplace:
        :param neg_slope:
        :param norm_affine:
        :param conv_bias:
        '''
        self.input_channels = input_feature_channels
        self.output_channels = output_feature_channels

        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}
        if conv_kwargs is None:
            conv_kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1, 'dilation': 1, 'bias': True}

        self.nonlin_kwargs = nonlin_kwargs
        self.nonlin = nonlin
        self.dropout_op = dropout_op
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.conv_kwargs = conv_kwargs
        self.conv_op = conv_op
        self.norm_op = norm_op

        if first_stride is not None:
            self.conv_kwargs_first_conv = deepcopy(conv_kwargs)
            self.conv_kwargs_first_conv['stride'] = first_stride
        else:
            self.conv_kwargs_first_conv = conv_kwargs

        super(StackedConvLayers, self).__init__()
        self.blocks = nn.Sequential(
            *([basic_block(input_feature_channels, output_feature_channels, self.conv_op,
                           self.conv_kwargs_first_conv,
                           self.norm_op, self.norm_op_kwargs, self.dropout_op, self.dropout_op_kwargs,
                           self.nonlin, self.nonlin_kwargs)] +
              [basic_block(output_feature_channels, output_feature_channels, self.conv_op,
                           self.conv_kwargs,
                           self.norm_op, self.norm_op_kwargs, self.dropout_op, self.dropout_op_kwargs,
                           self.nonlin, self.nonlin_kwargs) for _ in range(num_convs - 1)]))

    def forward(self, x):
        return self.blocks(x)


def print_module_training_status(module):
    if isinstance(module, nn.Conv2d) or isinstance(module, nn.Conv3d) or isinstance(module, nn.Dropout3d) or \
            isinstance(module, nn.Dropout2d) or isinstance(module, nn.Dropout) or isinstance(module, nn.InstanceNorm3d) \
            or isinstance(module, nn.InstanceNorm2d) or isinstance(module, nn.InstanceNorm1d) \
            or isinstance(module, nn.BatchNorm2d) or isinstance(module, nn.BatchNorm3d) or isinstance(module,
                                                                                                      nn.BatchNorm1d):
        print(str(module), module.training)


class Upsample(nn.Module):
    def __init__(self, size=None, scale_factor=None, mode='nearest', align_corners=False):
        super(Upsample, self).__init__()
        self.align_corners = align_corners
        self.mode = mode
        self.scale_factor = scale_factor
        self.size = size

    def forward(self, x):
        return nn.functional.interpolate(x, size=self.size, scale_factor=self.scale_factor, mode=self.mode,
                                         align_corners=self.align_corners)


class Generic_UNet(nn.Module):
    DEFAULT_BATCH_SIZE_3D = 10
    DEFAULT_PATCH_SIZE_3D = (48, 48, 48)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(self, input_channels, base_num_features, num_classes, num_pool, num_conv_per_stage=2,
                 feat_map_mul_on_downscale=2, conv_op=nn.Conv2d,
                 norm_op=nn.BatchNorm2d, norm_op_kwargs=None,
                 dropout_op=nn.Dropout2d, dropout_op_kwargs=None,
                 nonlin=nn.LeakyReLU, nonlin_kwargs=None, deep_supervision=True, dropout_in_localization=False,
                 final_nonlin=softmax_helper, weightInitializer=InitWeights_He(1e-2), pool_op_kernel_sizes=None,
                 conv_kernel_sizes=None,
                 upscale_logits=False, convolutional_pooling=False, convolutional_upsampling=False,
                 max_num_features=None, basic_block=ConvDropoutNormNonlin,
                 seg_output_use_bias=False):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super(Generic_UNet, self).__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            self.conv_blocks_context.append(StackedConvLayers(input_features, output_features, num_conv_per_stage,
                                                              self.conv_op, self.conv_kwargs, self.norm_op,
                                                              self.norm_op_kwargs, self.dropout_op,
                                                              self.dropout_op_kwargs, self.nonlin, self.nonlin_kwargs,
                                                              first_stride, basic_block=basic_block))
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]
        self.conv_blocks_context.append(nn.Sequential(
            StackedConvLayers(input_features, output_features, num_conv_per_stage - 1, self.conv_op, self.conv_kwargs,
                              self.norm_op, self.norm_op_kwargs, self.dropout_op, self.dropout_op_kwargs, self.nonlin,
                              self.nonlin_kwargs, first_stride, basic_block=basic_block),
            StackedConvLayers(output_features, final_num_features, 1, self.conv_op, self.conv_kwargs,
                              self.norm_op, self.norm_op_kwargs, self.dropout_op, self.dropout_op_kwargs, self.nonlin,
                              self.nonlin_kwargs, basic_block=basic_block)))

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(transpconv(nfeatures_from_down, nfeatures_from_skip, pool_op_kernel_sizes[-(u + 1)],
                                          pool_op_kernel_sizes[-(u + 1)], bias=False))

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[- (u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[- (u + 1)]
            self.conv_blocks_localization.append(nn.Sequential(
                StackedConvLayers(n_features_after_tu_and_concat, nfeatures_from_skip, num_conv_per_stage - 1,
                                  self.conv_op, self.conv_kwargs, self.norm_op, self.norm_op_kwargs, self.dropout_op,
                                  self.dropout_op_kwargs, self.nonlin, self.nonlin_kwargs, basic_block=basic_block),
                StackedConvLayers(nfeatures_from_skip, final_num_features, 1, self.conv_op, self.conv_kwargs,
                                  self.norm_op, self.norm_op_kwargs, self.dropout_op, self.dropout_op_kwargs,
                                  self.nonlin, self.nonlin_kwargs, basic_block=basic_block)
            ))

        if self.do_ds:
            for ds in range(len(self.conv_blocks_localization)):
                self.seg_outputs.append(conv_op(self.conv_blocks_localization[ds][-1].output_channels, num_classes,
                                                1, 1, 0, 1, 1, seg_output_use_bias))
        else:
            self.seg_outputs.append(conv_op(self.conv_blocks_localization[-1][-1].output_channels, num_classes,
                                            1, 1, 0, 1, 1, seg_output_use_bias))

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]),
                                                        mode=upsample_mode))
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops)  # lambda x:x is not a Module so we need to distinguish here

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)
            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        x = self.conv_blocks_context[-1](x)

        if self.do_ds:
            for u in range(len(self.tu)):
                x = self.tu[u](x)
                x = torch.cat((x, skips[-(u + 1)]), dim=1)
                x = self.conv_blocks_localization[u](x)
                seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))
        else:
            for u in range(len(self.tu)):
                x = self.tu[u](x)
                x = torch.cat((x, skips[-(u + 1)]), dim=1)
                x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[0](x)))

        if self._deep_supervision and self.do_ds:
            return tuple([seg_outputs[-1]] + [i(j) for i, j in
                                              zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])])
        else:
            return seg_outputs[-1]
            # return xt

    def functional_forward_trmode_3ds(self, xnor, weights):
        '''
        This is for the non-deepsupervision.
        I have 60 parameters..downsampling == 3
        also note that this is the old version, I may not use it in the future.
        '''
        listparameters = list(weights.items())
        nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': False}
        norm_op_kwargs = {'eps': 1e-5}
        conv_kwargs = {'stride': 1, 'padding': 1, 'dilation': 1}
        conv_kwargs_stride = {'stride': 2, 'padding': 1, 'dilation': 1}
        out = xnor
        skips = []
        seg_outputs = []

        for k in range(0, 2) :
            out = F.conv3d(out, weight=listparameters[24 + k*4][1], bias=listparameters[25+ k*4][1], **conv_kwargs)
            out = F.instance_norm(out, running_mean=None, running_var=None, weight=listparameters[26+ k*4][1],
                               bias=listparameters[27+ k*4][1], **norm_op_kwargs)
            out = F.leaky_relu(out, **nonlin_kwargs)
        skips.append(out)
        for k in range(2, 4) :
            if k == 2:
                out = F.conv3d(out, weight=listparameters[24 + k * 4][1], bias=listparameters[25 + k * 4][1],
                               **conv_kwargs_stride)
            else:
                out = F.conv3d(out, weight=listparameters[24 + k*4][1], bias=listparameters[25+ k*4][1], **conv_kwargs)
            out = F.instance_norm(out, running_mean=None, running_var=None, weight=listparameters[26+ k*4][1],
                               bias=listparameters[27+ k*4][1], **norm_op_kwargs)
            out = F.leaky_relu(out, **nonlin_kwargs)
        skips.append(out)
        for k in range(4, 6) :
            if k == 4:
                out = F.conv3d(out, weight=listparameters[24 + k * 4][1], bias=listparameters[25 + k * 4][1],
                               **conv_kwargs_stride)
            else:
                out = F.conv3d(out, weight=listparameters[24 + k * 4][1], bias=listparameters[25 + k * 4][1],
                               **conv_kwargs)
            out = F.instance_norm(out, running_mean=None, running_var=None, weight=listparameters[26+ k*4][1],
                               bias=listparameters[27+ k*4][1], **norm_op_kwargs)
            out = F.leaky_relu(out, **nonlin_kwargs)
        skips.append(out)
        for k in range(6, 8) :
            if k == 6:
                out = F.conv3d(out, weight=listparameters[24 + k * 4][1], bias=listparameters[25 + k * 4][1],
                               **conv_kwargs_stride)
            else:
                out = F.conv3d(out, weight=listparameters[24 + k * 4][1], bias=listparameters[25 + k * 4][1],
                               **conv_kwargs)
            out = F.instance_norm(out, running_mean=None, running_var=None, weight=listparameters[26+ k*4][1],
                               bias=listparameters[27+ k*4][1], **norm_op_kwargs)
            out = F.leaky_relu(out, **nonlin_kwargs)

        out = F.conv_transpose3d(out, weight = listparameters[56][1], stride = 2)
        out = torch.cat((out, skips[-1]), dim=1)
        for k in range(0, 2) :
            out = F.conv3d(out, weight=listparameters[0 + k*4][1], bias=listparameters[1+ k*4][1], **conv_kwargs)
            out = F.instance_norm(out, running_mean=None, running_var=None, weight=listparameters[2+ k*4][1],
                               bias=listparameters[3+ k*4][1], **norm_op_kwargs)
            out = F.leaky_relu(out, **nonlin_kwargs)
        out = F.conv_transpose3d(out, weight=listparameters[57][1], stride=2)
        out = torch.cat((out, skips[-2]), dim=1)
        for k in range(2, 4) :
            out = F.conv3d(out, weight=listparameters[0 + k*4][1], bias=listparameters[1+ k*4][1], **conv_kwargs)
            out = F.instance_norm(out, running_mean=None, running_var=None, weight=listparameters[2+ k*4][1],
                               bias=listparameters[3+ k*4][1], **norm_op_kwargs)
            out = F.leaky_relu(out, **nonlin_kwargs)
        out = F.conv_transpose3d(out, weight=listparameters[58][1], stride=2)
        out = torch.cat((out, skips[-3]), dim=1)
        for k in range(4, 6) :
            out = F.conv3d(out, weight=listparameters[0 + k*4][1], bias=listparameters[1+ k*4][1], **conv_kwargs)
            out = F.instance_norm(out, running_mean=None, running_var=None, weight=listparameters[2+ k*4][1],
                               bias=listparameters[3+ k*4][1], **norm_op_kwargs)
            out = F.leaky_relu(out, **nonlin_kwargs)

        seg_outputs.append(self.final_nonlin(F.conv3d(out, weight=listparameters[-1][1])))

        return seg_outputs[-1]

    def functional_forward_trmode(self, xnor, weights):
        '''
        This is for the deepsupervision version, and set downsampling to 4.
        Use the evaluation mode, which means use the given statistics.
        downsampling 4 has 80 parameters 
        '''
        listparameters = list(weights.items())
        nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': False}
        norm_op_kwargs = {'eps': 1e-5}
        conv_kwargs = {'stride': 1, 'padding': 1, 'dilation': 1}
        conv_kwargs_stride = {'stride': 2, 'padding': 1, 'dilation': 1}
        out = xnor
        skips = []
        seg_outputs = []

        for k in range(0, 2) :
            out = F.conv3d(out, weight=listparameters[32 + k*4][1], bias=listparameters[33+ k*4][1], **conv_kwargs)
            out = F.instance_norm(out, running_mean=None, running_var=None, weight=listparameters[34+ k*4][1],
                               bias=listparameters[35+ k*4][1], **norm_op_kwargs)
            out = F.leaky_relu(out, **nonlin_kwargs)
        skips.append(out)
        for k in range(2, 4) :
            if k == 2:
                out = F.conv3d(out, weight=listparameters[32 + k * 4][1], bias=listparameters[33 + k * 4][1],
                               **conv_kwargs_stride)
            else:
                out = F.conv3d(out, weight=listparameters[32 + k*4][1], bias=listparameters[33+ k*4][1], **conv_kwargs)
            out = F.instance_norm(out, running_mean=None, running_var=None, weight=listparameters[34+ k*4][1],
                               bias=listparameters[35+ k*4][1], **norm_op_kwargs)
            out = F.leaky_relu(out, **nonlin_kwargs)
        skips.append(out)
        for k in range(4, 6) :
            if k == 4:
                out = F.conv3d(out, weight=listparameters[32 + k * 4][1], bias=listparameters[33 + k * 4][1],
                               **conv_kwargs_stride)
            else:
                out = F.conv3d(out, weight=listparameters[32 + k * 4][1], bias=listparameters[33 + k * 4][1],
                               **conv_kwargs)
            out = F.instance_norm(out, running_mean=None, running_var=None, weight=listparameters[34+ k*4][1],
                               bias=listparameters[35+ k*4][1], **norm_op_kwargs)
            out = F.leaky_relu(out, **nonlin_kwargs)
        skips.append(out)
        for k in range(6, 8) :
            if k == 6:
                out = F.conv3d(out, weight=listparameters[32 + k * 4][1], bias=listparameters[33 + k * 4][1],
                               **conv_kwargs_stride)
            else:
                out = F.conv3d(out, weight=listparameters[32 + k * 4][1], bias=listparameters[33 + k * 4][1],
                               **conv_kwargs)
            out = F.instance_norm(out, running_mean=None, running_var=None, weight=listparameters[34+ k*4][1],
                               bias=listparameters[35+ k*4][1], **norm_op_kwargs)
            out = F.leaky_relu(out, **nonlin_kwargs)
        skips.append(out)
        for k in range(8, 10) :
            if k == 8:
                out = F.conv3d(out, weight=listparameters[32 + k * 4][1], bias=listparameters[33 + k * 4][1],
                               **conv_kwargs_stride)
            else:
                out = F.conv3d(out, weight=listparameters[32 + k * 4][1], bias=listparameters[33 + k * 4][1],
                               **conv_kwargs)
            out = F.instance_norm(out, running_mean=None, running_var=None, weight=listparameters[34+ k*4][1],
                               bias=listparameters[35+ k*4][1], **norm_op_kwargs)
            out = F.leaky_relu(out, **nonlin_kwargs)
        
        out = F.conv_transpose3d(out, weight = listparameters[72][1], stride = 2)
        out = torch.cat((out, skips[-1]), dim=1)
        for k in range(0, 2) :
            out = F.conv3d(out, weight=listparameters[0 + k*4][1], bias=listparameters[1+ k*4][1], **conv_kwargs)
            out = F.instance_norm(out, running_mean=None, running_var=None, weight=listparameters[2+ k*4][1],
                               bias=listparameters[3+ k*4][1], **norm_op_kwargs)
            out = F.leaky_relu(out, **nonlin_kwargs)
        seg_outputs.append(self.final_nonlin(F.conv3d(out, weight=listparameters[-4][1])))
        out = F.conv_transpose3d(out, weight=listparameters[73][1], stride=2)
        out = torch.cat((out, skips[-2]), dim=1)
        for k in range(2, 4) :
            out = F.conv3d(out, weight=listparameters[0 + k*4][1], bias=listparameters[1+ k*4][1], **conv_kwargs)
            out = F.instance_norm(out, running_mean=None, running_var=None, weight=listparameters[2+ k*4][1],
                               bias=listparameters[3+ k*4][1], **norm_op_kwargs)
            out = F.leaky_relu(out, **nonlin_kwargs)
        seg_outputs.append(self.final_nonlin(F.conv3d(out, weight=listparameters[-3][1])))
        out = F.conv_transpose3d(out, weight=listparameters[74][1], stride=2)
        out = torch.cat((out, skips[-3]), dim=1)
        for k in range(4, 6) :
            out = F.conv3d(out, weight=listparameters[0 + k*4][1], bias=listparameters[1+ k*4][1], **conv_kwargs)
            out = F.instance_norm(out, running_mean=None, running_var=None, weight=listparameters[2+ k*4][1],
                               bias=listparameters[3+ k*4][1], **norm_op_kwargs)
            out = F.leaky_relu(out, **nonlin_kwargs)
        seg_outputs.append(self.final_nonlin(F.conv3d(out, weight=listparameters[-2][1])))
        out = F.conv_transpose3d(out, weight=listparameters[75][1], stride=2)
        out = torch.cat((out, skips[-4]), dim=1)
        for k in range(6, 8) :
            out = F.conv3d(out, weight=listparameters[0 + k*4][1], bias=listparameters[1+ k*4][1], **conv_kwargs)
            out = F.instance_norm(out, running_mean=None, running_var=None, weight=listparameters[2+ k*4][1],
                               bias=listparameters[3+ k*4][1], **norm_op_kwargs)
            out = F.leaky_relu(out, **nonlin_kwargs)

        seg_outputs.append(self.final_nonlin(F.conv3d(out, weight=listparameters[-1][1])))

        return tuple([seg_outputs[-1]] + [i for i in seg_outputs[:-1][::-1]])

    def functional_forward_trmode_returnstats(self, xnor, weights):
        '''
        This is for the deepsupervision version, and set downsampling to 4.
        Use the evaluation mode, which means use the given statistics.
        downsampling 4 has 80 parameters 
        '''
        from testIN import ForwardIN

        meanlist = []
        varlist = []

        listparameters = list(weights.items())
        nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': False}
        norm_op_kwargs = {'eps': 1e-5}
        conv_kwargs = {'stride': 1, 'padding': 1, 'dilation': 1}
        conv_kwargs_stride = {'stride': 2, 'padding': 1, 'dilation': 1}
        out = xnor
        skips = []
        seg_outputs = []

        for k in range(0, 2) :
            out = F.conv3d(out, weight=listparameters[32 + k*4][1], bias=listparameters[33+ k*4][1], **conv_kwargs)
            out, mean, var = ForwardIN(out, weight=listparameters[34+ k*4][1], bias=listparameters[35+ k*4][1])
            meanlist.append(mean)
            varlist.append(var)
            out = F.leaky_relu(out, **nonlin_kwargs)
        skips.append(out)
        for k in range(2, 4) :
            if k == 2:
                out = F.conv3d(out, weight=listparameters[32 + k * 4][1], bias=listparameters[33 + k * 4][1],
                            **conv_kwargs_stride)
            else:
                out = F.conv3d(out, weight=listparameters[32 + k*4][1], bias=listparameters[33+ k*4][1], **conv_kwargs)
            out, mean, var = ForwardIN(out, weight=listparameters[34+ k*4][1], bias=listparameters[35+ k*4][1])
            meanlist.append(mean)
            varlist.append(var)
            out = F.leaky_relu(out, **nonlin_kwargs)
        skips.append(out)
        for k in range(4, 6) :
            if k == 4:
                out = F.conv3d(out, weight=listparameters[32 + k * 4][1], bias=listparameters[33 + k * 4][1],
                            **conv_kwargs_stride)
            else:
                out = F.conv3d(out, weight=listparameters[32 + k * 4][1], bias=listparameters[33 + k * 4][1],
                            **conv_kwargs)
            out, mean, var = ForwardIN(out, weight=listparameters[34+ k*4][1], bias=listparameters[35+ k*4][1])
            meanlist.append(mean)
            varlist.append(var)
            out = F.leaky_relu(out, **nonlin_kwargs)
        skips.append(out)
        for k in range(6, 8) :
            if k == 6:
                out = F.conv3d(out, weight=listparameters[32 + k * 4][1], bias=listparameters[33 + k * 4][1],
                            **conv_kwargs_stride)
            else:
                out = F.conv3d(out, weight=listparameters[32 + k * 4][1], bias=listparameters[33 + k * 4][1],
                            **conv_kwargs)
            out, mean, var = ForwardIN(out, weight=listparameters[34+ k*4][1], bias=listparameters[35+ k*4][1])
            meanlist.append(mean)
            varlist.append(var)
            out = F.leaky_relu(out, **nonlin_kwargs)
        skips.append(out)
        for k in range(8, 10) :
            if k == 8:
                out = F.conv3d(out, weight=listparameters[32 + k * 4][1], bias=listparameters[33 + k * 4][1],
                            **conv_kwargs_stride)
            else:
                out = F.conv3d(out, weight=listparameters[32 + k * 4][1], bias=listparameters[33 + k * 4][1],
                            **conv_kwargs)
            out, mean, var = ForwardIN(out, weight=listparameters[34+ k*4][1], bias=listparameters[35+ k*4][1])
            meanlist.append(mean)
            varlist.append(var)
            out = F.leaky_relu(out, **nonlin_kwargs)
        
        out = F.conv_transpose3d(out, weight = listparameters[72][1], stride = 2)
        out = torch.cat((out, skips[-1]), dim=1)
        for k in range(0, 2) :
            out = F.conv3d(out, weight=listparameters[0 + k*4][1], bias=listparameters[1+ k*4][1], **conv_kwargs)
            out, mean, var = ForwardIN(out, weight=listparameters[2+ k*4][1], bias=listparameters[3+ k*4][1])
            meanlist.append(mean)
            varlist.append(var)
            out = F.leaky_relu(out, **nonlin_kwargs)
        seg_outputs.append(self.final_nonlin(F.conv3d(out, weight=listparameters[-4][1])))
        out = F.conv_transpose3d(out, weight=listparameters[73][1], stride=2)
        out = torch.cat((out, skips[-2]), dim=1)
        for k in range(2, 4) :
            out = F.conv3d(out, weight=listparameters[0 + k*4][1], bias=listparameters[1+ k*4][1], **conv_kwargs)
            out, mean, var = ForwardIN(out, weight=listparameters[2+ k*4][1], bias=listparameters[3+ k*4][1])
            meanlist.append(mean)
            varlist.append(var)
            out = F.leaky_relu(out, **nonlin_kwargs)
        seg_outputs.append(self.final_nonlin(F.conv3d(out, weight=listparameters[-3][1])))
        out = F.conv_transpose3d(out, weight=listparameters[74][1], stride=2)
        out = torch.cat((out, skips[-3]), dim=1)
        for k in range(4, 6) :
            out = F.conv3d(out, weight=listparameters[0 + k*4][1], bias=listparameters[1+ k*4][1], **conv_kwargs)
            out, mean, var = ForwardIN(out, weight=listparameters[2+ k*4][1], bias=listparameters[3+ k*4][1])
            meanlist.append(mean)
            varlist.append(var)
            out = F.leaky_relu(out, **nonlin_kwargs)
        seg_outputs.append(self.final_nonlin(F.conv3d(out, weight=listparameters[-2][1])))
        out = F.conv_transpose3d(out, weight=listparameters[75][1], stride=2)
        out = torch.cat((out, skips[-4]), dim=1)
        for k in range(6, 8) :
            out = F.conv3d(out, weight=listparameters[0 + k*4][1], bias=listparameters[1+ k*4][1], **conv_kwargs)
            out, mean, var = ForwardIN(out, weight=listparameters[2+ k*4][1], bias=listparameters[3+ k*4][1])
            meanlist.append(mean)
            varlist.append(var)
            out = F.leaky_relu(out, **nonlin_kwargs)

        seg_outputs.append(self.final_nonlin(F.conv3d(out, weight=listparameters[-1][1])))

        return tuple([seg_outputs[-1]] + [i for i in seg_outputs[:-1][::-1]]), meanlist, varlist

    def functional_forward_evalmode(self, xnor, weights, meanlist, varlist, maskinit=None):
        '''
        This is for the deepsupervision version, and set downsampling to 4.
        Use the evaluation mode, which means use the given statistics.
        downsampling 4 has 80 parameters 
        '''

        from testIN import ForwardIN
        
        listparameters = list(weights.items())
        nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': False}
        norm_op_kwargs = {'eps': 1e-5}
        conv_kwargs = {'stride': 1, 'padding': 1, 'dilation': 1}
        conv_kwargs_stride = {'stride': 2, 'padding': 1, 'dilation': 1}
        out = xnor
        skips = []
        seg_outputs = []

        bncount = 0

        for k in range(0, 2) :
            out = F.conv3d(out, weight=listparameters[32 + k*4][1], **conv_kwargs)
            out = out + listparameters[33+ k*4][1][None, :, None, None, None]
            out, _, _ = ForwardIN(out, weight=listparameters[34+ k*4][1], bias=listparameters[35+ k*4][1], trmode = False, 
                            running_mean=meanlist[bncount], running_var=varlist[bncount])
            bncount += 1
            out = F.leaky_relu(out, **nonlin_kwargs)
        skips.append(out)
        for k in range(2, 4) :
            if k == 2:
                out = F.conv3d(out, weight=listparameters[32 + k*4][1], **conv_kwargs_stride)
                out = out + listparameters[33+ k*4][1][None, :, None, None, None]
            else:
                out = F.conv3d(out, weight=listparameters[32 + k*4][1], **conv_kwargs)
                out = out + listparameters[33+ k*4][1][None, :, None, None, None]
            out, _, _ = ForwardIN(out, weight=listparameters[34+ k*4][1], bias=listparameters[35+ k*4][1], trmode = False, 
                            running_mean=meanlist[bncount], running_var=varlist[bncount])
            bncount += 1
            out = F.leaky_relu(out, **nonlin_kwargs)
        skips.append(out)
        for k in range(4, 6) :
            if k == 4:
                out = F.conv3d(out, weight=listparameters[32 + k*4][1], **conv_kwargs_stride)
                out = out + listparameters[33+ k*4][1][None, :, None, None, None]
            else:
                out = F.conv3d(out, weight=listparameters[32 + k*4][1], **conv_kwargs)
                out = out + listparameters[33+ k*4][1][None, :, None, None, None]
            out, _, _ = ForwardIN(out, weight=listparameters[34+ k*4][1], bias=listparameters[35+ k*4][1], trmode = False, 
                            running_mean=meanlist[bncount], running_var=varlist[bncount])
            bncount += 1
            out = F.leaky_relu(out, **nonlin_kwargs)
        skips.append(out)
        for k in range(6, 8) :
            if k == 6:
                out = F.conv3d(out, weight=listparameters[32 + k*4][1], **conv_kwargs_stride)
                out = out + listparameters[33+ k*4][1][None, :, None, None, None]
            else:
                out = F.conv3d(out, weight=listparameters[32 + k*4][1], **conv_kwargs)
                out = out + listparameters[33+ k*4][1][None, :, None, None, None]
            out, _, _ = ForwardIN(out, weight=listparameters[34+ k*4][1], bias=listparameters[35+ k*4][1], trmode = False, 
                            running_mean=meanlist[bncount], running_var=varlist[bncount])
            bncount += 1
            out = F.leaky_relu(out, **nonlin_kwargs)
        skips.append(out)
        for k in range(8, 10) :
            if k == 8:
                out = F.conv3d(out, weight=listparameters[32 + k*4][1], **conv_kwargs_stride)
                out = out + listparameters[33+ k*4][1][None, :, None, None, None]
            else:
                out = F.conv3d(out, weight=listparameters[32 + k*4][1], **conv_kwargs)
                out = out + listparameters[33+ k*4][1][None, :, None, None, None]
            out, _, _ = ForwardIN(out, weight=listparameters[34+ k*4][1], bias=listparameters[35+ k*4][1], trmode = False, 
                            running_mean=meanlist[bncount], running_var=varlist[bncount])
            bncount += 1
            out = F.leaky_relu(out, **nonlin_kwargs)
        
        out = F.conv_transpose3d(out, weight = listparameters[72][1], stride = 2)
        out = torch.cat((out, skips[-1]), dim=1)
        for k in range(0, 2) :
            out = F.conv3d(out, weight=listparameters[0 + k*4][1], **conv_kwargs)
            out = out + listparameters[1+ k*4][1][None, :, None, None, None]
            out, _, _ = ForwardIN(out, weight=listparameters[2+ k*4][1], bias=listparameters[3+ k*4][1], trmode = False, 
                            running_mean=meanlist[bncount], running_var=varlist[bncount])
            bncount += 1
            out = F.leaky_relu(out, **nonlin_kwargs)
        seg_outputs.append(self.final_nonlin(F.conv3d(out, weight=listparameters[-4][1])))
        out = F.conv_transpose3d(out, weight=listparameters[73][1], stride=2)
        out = torch.cat((out, skips[-2]), dim=1)
        for k in range(2, 4) :
            out = F.conv3d(out, weight=listparameters[0 + k*4][1], **conv_kwargs)
            out = out + listparameters[1+ k*4][1][None, :, None, None, None]
            out, _, _ = ForwardIN(out, weight=listparameters[2+ k*4][1], bias=listparameters[3+ k*4][1], trmode = False, 
                            running_mean=meanlist[bncount], running_var=varlist[bncount])
            bncount += 1
            out = F.leaky_relu(out, **nonlin_kwargs)
        seg_outputs.append(self.final_nonlin(F.conv3d(out, weight=listparameters[-3][1])))
        out = F.conv_transpose3d(out, weight=listparameters[74][1], stride=2)
        out = torch.cat((out, skips[-3]), dim=1)
        for k in range(4, 6) :
            out = F.conv3d(out, weight=listparameters[0 + k*4][1], **conv_kwargs)
            out = out + listparameters[1+ k*4][1][None, :, None, None, None]
            out, _, _ = ForwardIN(out, weight=listparameters[2+ k*4][1], bias=listparameters[3+ k*4][1], trmode = False, 
                            running_mean=meanlist[bncount], running_var=varlist[bncount])
            bncount += 1
            out = F.leaky_relu(out, **nonlin_kwargs)
        seg_outputs.append(self.final_nonlin(F.conv3d(out, weight=listparameters[-2][1])))
        out = F.conv_transpose3d(out, weight=listparameters[75][1], stride=2)
        out = torch.cat((out, skips[-4]), dim=1)
        for k in range(6, 8) :
            out = F.conv3d(out, weight=listparameters[0 + k*4][1], **conv_kwargs)
            out = out + listparameters[1+ k*4][1][None, :, None, None, None]
            out, _, _ = ForwardIN(out, weight=listparameters[2+ k*4][1], bias=listparameters[3+ k*4][1], trmode = False, 
                            running_mean=meanlist[bncount], running_var=varlist[bncount])
            bncount += 1
            out = F.leaky_relu(out, **nonlin_kwargs)

        seg_outputs.append(self.final_nonlin(F.conv3d(out, weight=listparameters[-1][1])))

        return tuple([seg_outputs[-1]] + [i for i in seg_outputs[:-1][::-1]])