ml_modules.py

import torch
from torch import nn
import torch.fft
import numpy as np
import math
import functools
import warnings

from math import sqrt


# grab from upstream pytorch branch and paste here for now
def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1. + math.erf(x / math.sqrt(2.))) / 2.

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
                      "The distribution of values may be incorrect.",
                      stacklevel=2)

    with torch.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        # Get upper and lower cdf values
        l = norm_cdf((a - mean) / std)
        u = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [l, u], then translate to
        # [2l-1, 2u-1].
        tensor.uniform_(2 * l - 1, 2 * u - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor.erfinv_()

        # Transform to proper mean, std
        tensor.mul_(std * math.sqrt(2.))
        tensor.add_(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp_(min=a, max=b)
        return tensor


def init_weights_requ(m):
    if type(m) == nn.Linear:
        if hasattr(m, 'weight'):
            nn.init.kaiming_normal_(m.weight, a=0.0, nonlinearity='relu', mode='fan_in')
        # if hasattr(m, 'bias'):
        #     nn.init.uniform_(m.bias, -.5,.5)


def init_weights_trunc_normal(m):
    if type(m) == nn.Linear:
        if hasattr(m, 'weight'):
            fan_in = m.weight.size(1)
            fan_out = m.weight.size(0)
            std = math.sqrt(2.0 / float(fan_in + fan_out))
            mean = 0.
            # initialize with the same behavior as tf.truncated_normal
            # "The generated values follow a normal distribution with specified mean and
            # standard deviation, except that values whose magnitude is more than 2
            # standard deviations from the mean are dropped and re-picked."
            _no_grad_trunc_normal_(m.weight, mean, std, -2 * std, 2 * std)


# TODO: changed initialization for backward graph 'fan out'
def init_weights_normal(m):
    if type(m) == nn.Linear:
        if hasattr(m, 'weight'):
            nn.init.kaiming_normal_(m.weight, a=0.0, nonlinearity='relu', mode='fan_out')
        if hasattr(m, 'bias'):
            nn.init.uniform_(m.bias, -1, 1)
            # m.bias.data.fill_(0.)


def init_weights_selu(m):
    if type(m) == nn.Linear:
        if hasattr(m, 'weight'):
            num_input = m.weight.size(-1)
            nn.init.normal_(m.weight, std=1 / math.sqrt(num_input))
        # if hasattr(m, 'bias'):
        #     m.bias.data.fill_(0.)


def init_weights_elu(m):
    if type(m) == nn.Linear:
        if hasattr(m, 'weight'):
            num_input = m.weight.size(-1)
            nn.init.normal_(m.weight, std=math.sqrt(1.5505188080679277) / math.sqrt(num_input))
        # if hasattr(m, 'bias'):
        #     m.bias.data.fill_(0.)


def init_weights_xavier(m):
    if type(m) == nn.Linear:
        if hasattr(m, 'weight'):
            nn.init.xavier_normal_(m.weight)
        if hasattr(m, 'bias'):
            m.bias.data.fill_(0.)


def init_weights_uniform(m):
    if type(m) == nn.Linear:
        if hasattr(m, 'weight'):
            torch.nn.init.kaiming_uniform_(m.weight, nonlinearity='relu')


def sine_init(m, w0=30):
    with torch.no_grad():
        if hasattr(m, 'weight'):
            num_input = m.weight.size(-1)
            m.weight.uniform_(-np.sqrt(6 / num_input) / w0, np.sqrt(6 / num_input) / w0)


def first_layer_sine_init(m):
    with torch.no_grad():
        if hasattr(m, 'weight'):
            num_input = m.weight.size(-1)
            m.weight.uniform_(-1 / num_input, 1 / num_input)


def first_layer_sine_wavelet_init(m):
    with torch.no_grad():
        if hasattr(m, 'weight'):
            num_input = m.weight.size(-1)
            # See paper sec. 3.2, final paragraph, and supplement Sec. 1.5 for discussion of factor 30
            m.weight.uniform_(-np.sqrt(6 / num_input), np.sqrt(6 / num_input))
            m.bias.uniform_(-1, 1)


def increase_w0(model, factor=1.1, max=20):
    with torch.no_grad():
        linear = []
        sine = []
        for m in model.modules():
            if hasattr(m, 'weight'):
                linear.append(m)
            elif type(m) == Sine or type(m) == FirstSine:
                if hasattr(m, 'w0'):
                    sine.append(m)

        # remember there can be one more linear layer than the # of sine layers
        for idx, m in enumerate(sine):
            w0_prev = m.w0.clone()
            m.w0 = torch.clamp(m.w0 * factor, 0, max)

            applied_factor = m.w0 / w0_prev
            linear[idx].weight.data = linear[idx].weight.data / applied_factor
            linear[idx].bias.data = linear[idx].bias.data / applied_factor

        print(sine[0].w0)


def compl_conj(x):
    y = x.clone()
    y[..., 1::2] = -1 * y[..., 1::2]
    return y


def compl_div(x, y):
    ''' x / y '''
    a = x[..., ::2]
    b = x[..., 1::2]
    c = y[..., ::2]
    d = y[..., 1::2]

    outr = (a * c + b * d) / (c ** 2 + d ** 2)
    outi = (b * c - a * d) / (c ** 2 + d ** 2)
    out = torch.zeros_like(x)
    out[..., ::2] = outr
    out[..., 1::2] = outi
    return out
    # return torch.cat((outr[..., None], outi[..., None]), dim=-1)


def compl_mul(x, y):
    '''  x * y '''
    a = x[..., ::2]
    b = x[..., 1::2]
    c = y[..., ::2]
    d = y[..., 1::2]

    outr = a * c - b * d
    outi = (a + b) * (c + d) - a * c - b * d
    out = torch.zeros_like(x)
    out[..., ::2] = outr
    out[..., 1::2] = outi
    return out
    # return torch.cat((outr[..., None], outi[..., None]), dim=-1)


class FirstSine(nn.Module):
    def __init__(self, w0=20):
        super().__init__()
        self.w0 = torch.tensor(w0)

    def forward(self, input):
        return torch.sin(self.w0 * input)


class Sine(nn.Module):
    def __init__(self, w0=20):
        super().__init__()
        self.w0 = torch.tensor(w0)

    def forward(self, input):
        # return torch.sin(np.sqrt(256)*input)
        return torch.sin(self.w0 * input)


class SineWavelet(nn.Module):
    def __init__(self, w0=20, num_features=256):
        super().__init__()
        self.w0 = torch.tensor(w0)
        self.sigma = nn.Parameter(torch.rand(num_features, dtype=torch.float32),
                                  requires_grad=True)

    def forward(self, input):
        # return torch.sin(np.sqrt(256)*input)
        # return torch.sin(self.w0*input)*torch.exp(-self.sigma * input**2)
        return torch.sin(self.w0 * input) * torch.relu(1. - self.sigma * torch.abs(input))


class RandSine(nn.Module):
    def __init__(self, mu_w0=50, std_w0=40, num_features=256):  # 30, 29
        super().__init__()
        self.w0 = mu_w0 + 2. * std_w0 * (torch.rand(num_features, dtype=torch.float32) - .5)

    def forward(self, input):
        # See paper sec. 3.2, final paragraph, and supplement Sec. 1.5 for discussion of factor 30
        return torch.sin(self.w0 * input)


class ReQU(nn.Module):
    def __init__(self, inplace=True):
        super().__init__()
        self.relu = nn.ReLU(inplace)

    def forward(self, input):
        # return torch.sin(np.sqrt(256)*input)
        return .5 * self.relu(input) ** 2


class MSoftplus(nn.Module):
    def __init__(self):
        super().__init__()
        self.softplus = nn.Softplus()
        self.cst = torch.log(torch.tensor(2.))

    def forward(self, input):
        return self.softplus(input) - self.cst


class Swish(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, input):
        return input * torch.sigmoid(input)


class ReQLU(nn.Module):
    def __init__(self):
        super().__init__()
        self.p_sq = 1 ** 2

    def forward(self, input):
        r_input = torch.relu(input)
        return self.p_sq * (torch.sqrt(1. + r_input ** 2 / self.p_sq) - 1.)

''' FCNet'''
def layer_factory(layer_type):
    layer_dict = \
        {'relu': (nn.ReLU(inplace=True), init_weights_normal),
         'lrelu': (nn.LeakyReLU(inplace=True), init_weights_normal),
         'requ': (ReQU(inplace=False), init_weights_requ),
         'reqlu': (ReQLU, init_weights_normal),
         'sigmoid': (nn.Sigmoid(), init_weights_xavier),
         'fsine': (Sine(), first_layer_sine_init),
         'sine': (Sine(), sine_init),
         'randsine': (RandSine(), sine_init),
         'tanh': (nn.Tanh(), init_weights_xavier),
         'htanh': (nn.Hardtanh(), init_weights_xavier),
         'ssign': (nn.Softsign(), init_weights_xavier),
         'selu': (nn.SELU(inplace=True), init_weights_selu),
         'gelu': (nn.GELU(), init_weights_selu),
         'silu': (nn.SiLU(), init_weights_selu),
         'swish': (Swish(), init_weights_selu),
         'softplus': (nn.Softplus(), init_weights_normal),
         'msoftplus': (MSoftplus(), init_weights_normal),
         'elu': (nn.ELU(), init_weights_elu)
         }
    return layer_dict[layer_type]

class EqualLR:
    def __init__(self, name):
        self.name = name

    def compute_weight(self, module):
        weight = getattr(module, self.name + '_orig')
        fan_in = weight.data.size(1) * weight.data[0][0].numel()

        return weight * sqrt(2 / fan_in)

    @staticmethod
    def apply(module, name):
        fn = EqualLR(name)

        weight = getattr(module, name)
        del module._parameters[name]
        module.register_parameter(name + '_orig', nn.Parameter(weight.data))
        module.register_forward_pre_hook(fn)

        return fn

    def __call__(self, module, input):
        weight = self.compute_weight(module)
        setattr(module, self.name, weight)


def equal_lr(module, name='weight'):
    EqualLR.apply(module, name)

    return module

class EqualConv2d(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__()

        conv = nn.Conv2d(*args, **kwargs)
        conv.weight.data.normal_()
        # torch.nn.init.kaiming_normal_(conv.weight)
        conv.bias.data.zero_()
        self.conv = equal_lr(conv)

    def forward(self, input):
        return self.conv(input)

class EqualLinear(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__()

        dense = nn.Linear(*args, **kwargs)
        # dense.weight.data.normal_()
        # dense.bias.data.zero_()
        self.dense = equal_lr(dense)

    def forward(self, input):
        return self.dense(input)

class FCBlock(nn.Module):
    '''A fully connected neural network that also allows swapping out the weights when used with a hypernetwork.
    Can be used just as a normal neural network though, as well.
    '''

    def __init__(self, in_features, features, out_features,
                 nonlinearity='relu', last_nonlinearity=None,
                 batch_norm=False, equalized=False, dropout=False):
        super().__init__()

        # Create hidden features list
        self.hidden_features = [int(in_features)]
        if features != []:
            self.hidden_features.extend(features)
        self.hidden_features.append(int(out_features))

        self.net = []
        for i in range(len(self.hidden_features) - 1):
            # Not the last
            hidden = False
            if i < len(self.hidden_features) - 2:
                nl = layer_factory(nonlinearity)[0]
                init = layer_factory(nonlinearity)[1]
                hidden = True
            # The last layer
            else:
                if last_nonlinearity is not None:
                    nl = layer_factory(last_nonlinearity)[0]
                    init = layer_factory(last_nonlinearity)[1]
            if equalized:
                layer = EqualLinear(self.hidden_features[i], self.hidden_features[i + 1])
            else:
                layer = nn.Linear(self.hidden_features[i], self.hidden_features[i + 1])
            if last_nonlinearity is not None:
                init(layer)
                self.net.append(layer)
                self.net.append(nl)
            else:
                # init_weights_normal(layer)
                self.net.append(layer)
            if hidden:
                if dropout:
                    self.net.append(nn.Dropout(p=0.5))
                if batch_norm:
                    self.net.append(nn.BatchNorm1d(num_features=self.hidden_features[i + 1]))

        self.net = nn.Sequential(*self.net)

    def forward(self, coords):
        output = self.net(coords)
        return output

''' PE '''
class PositionalEncoding(nn.Module):  # MetaModule):
    def __init__(self, num_encoding_functions=6, include_input=True, log_sampling=True, normalize=False,
                 input_dim=3, gaussian_pe=False, gaussian_variance=38):
        super().__init__()
        self.num_encoding_functions = num_encoding_functions
        self.include_input = include_input
        self.log_sampling = log_sampling
        self.normalize = normalize
        self.gaussian_pe = gaussian_pe
        self.normalization = None

        if self.gaussian_pe:
            # this needs to be registered as a parameter so that it is saved in the model state dict
            # and so that it is converted using .cuda(). Doesn't need to be trained though
            self.gaussian_weights = nn.Parameter(gaussian_variance * torch.randn(num_encoding_functions, input_dim),
                                                 requires_grad=False)

            # TODO: Normalization?

        else:
            self.frequency_bands = None
            if self.log_sampling:
                self.frequency_bands = 2.0 ** torch.linspace(
                    0.0,
                    self.num_encoding_functions - 1,
                    self.num_encoding_functions)
            else:
                self.frequency_bands = torch.linspace(
                    2.0 ** 0.0,
                    2.0 ** (self.num_encoding_functions - 1),
                    self.num_encoding_functions)

            if normalize:
                # self.normalization = nn.Parameter(1/self.frequency_bands.clone(), requires_grad=False)
                self.normalization = torch.tensor(1 / self.frequency_bands)

    def forward(self, tensor) -> torch.Tensor:
        r"""Apply positional encoding to the input.

        Args:
            tensor (torch.Tensor): Input tensor to be positionally encoded.
            encoding_size (optional, int): Number of encoding functions used to compute
                a positional encoding (default: 6).
            include_input (optional, bool): Whether or not to include the input in the
                positional encoding (default: True).

        Returns:
        (torch.Tensor): Positional encoding of the input tensor.
        """

        encoding = [tensor] if self.include_input else []
        if self.gaussian_pe:
            for func in [torch.sin, torch.cos]:
                encoding.append(func(torch.matmul(tensor, self.gaussian_weights.T)))
        else:
            for idx, freq in enumerate(self.frequency_bands):
                for func in [torch.sin, torch.cos]:
                    if self.normalization is not None:
                        encoding.append(self.normalization[idx] * func(tensor * freq))
                    else:
                        encoding.append(func(tensor * freq))

        # Special case, for no positional encoding
        if len(encoding) == 1:
            return encoding[0]
        else:
            return torch.cat(encoding, dim=-1)

''' SIREN '''
class SIREN(nn.Module):
    def __init__(self, in_features, out_features,
                 num_hidden_layers, hidden_features,
                 outermost_linear=False, w0=30):
        super(SIREN, self).__init__()

        nl = Sine(w0)
        first_nl = FirstSine(w0)
        self.weight_init = functools.partial(sine_init, w0=w0)
        self.first_layer_init = first_layer_sine_init

        self.net = []
        self.net.append(nn.Sequential(
            nn.Linear(in_features, hidden_features),
            first_nl
        ))

        for i in range(num_hidden_layers):
            self.net.append(nn.Sequential(
                nn.Linear(hidden_features, hidden_features),
                nl
            ))

        if outermost_linear:
            self.net.append(nn.Sequential(
                nn.Linear(hidden_features, out_features),
            ))
        else:
            self.net.append(nn.Sequential(
                nn.Linear(hidden_features, out_features),
                nl
            ))

        self.net = nn.Sequential(*self.net)
        if self.weight_init is not None:
            self.net.apply(self.weight_init)

        if self.first_layer_init is not None:
            self.net[0].apply(self.first_layer_init)

    def forward(self, coords):
        '''Simple forward pass without computation of spatial gradients.'''
        output = self.net(coords)
        return output

''' CNNs '''

'''Attention modules taken from https://github.com/luuuyi/CBAM.PyTorch'''
def conv3x3(in_planes, out_planes, stride=1, bias=True):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=1, bias=bias)
def deconv3x3(in_planes, out_planes, stride=1, bias=True):
    """3x3 seconvolution with padding"""
    return nn.ConvTranspose2d(in_planes, out_planes, kernel_size=3, padding=1,stride=stride, bias=bias)


class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

        self.fc = nn.Sequential(nn.Conv2d(in_planes, in_planes // 16, 1, bias=False),
                                nn.ReLU(),
                                nn.Conv2d(in_planes // 16, in_planes, 1, bias=False))
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.fc(self.avg_pool(x))
        max_out = self.fc(self.max_pool(x))
        out = avg_out + max_out
        return self.sigmoid(out)

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()

        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        x = torch.cat([avg_out, max_out], dim=1)
        x = self.conv1(x)
        return self.sigmoid(x)

class DoubleConvBlock(nn.Module):
    def __init__(self, in_size, out_size, padding, batch_norm,
                 equalized=False, dropout=False, attention=False):
        super(DoubleConvBlock, self).__init__()
        self.batch_norm = batch_norm
        self.dropout = dropout
        self.attention = attention
        if equalized:
            self.conv1 = EqualConv2d(in_size, out_size, kernel_size=3, padding=int(padding))
            self.conv2 = EqualConv2d(out_size, out_size, kernel_size=3, padding=int(padding))
        else:
            self.conv1 = conv3x3(in_size, out_size)
            self.conv2 = conv3x3(out_size, out_size)

        self.relu = nn.ReLU(inplace=True)
        if dropout:
            self.drop = nn.Dropout(p=0.2)
        if batch_norm:
            self.bn1 = nn.BatchNorm2d(out_size)
            self.bn2 = nn.BatchNorm2d(out_size)
        if attention:
            self.ca = ChannelAttention(out_size)
            self.sa = SpatialAttention()

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        if self.batch_norm:
            out = self.bn1(out)
        out = self.relu(out)
        if self.dropout:
            out = self.drop(out)

        out = self.conv2(out)
        if self.batch_norm:
            out = self.bn2(out)
        if self.attention:
            out = self.ca(out) * out
            out = self.sa(out) * out

        if self.dropout:
            out = self.drop(out)

        out = self.relu(out)

        return out
    
    
class DoubleDeconvBlock(nn.Module):
    def __init__(self, in_size, out_size, padding, batch_norm,
                 equalized=False, dropout=False, attention=False, last_relu=True):
        super(DoubleDeconvBlock, self).__init__()
        self.batch_norm = batch_norm
        self.dropout = dropout
        self.attention = attention
        self.last_relu=last_relu
        if equalized:
            self.conv1 = EqualDeconv2d(in_size, out_size, kernel_size=3, padding=int(padding))
            self.conv2 = EqualDeconv2d(out_size, out_size, kernel_size=3, padding=int(padding))
        else:
            self.conv1 = deconv3x3(in_size, out_size)
            self.conv2 = deconv3x3(out_size, out_size)

        self.relu = nn.ReLU(inplace=True)
        if dropout:
            self.drop = nn.Dropout(p=0.2)
        if batch_norm:
            self.bn1 = nn.BatchNorm2d(out_size)
            self.bn2 = nn.BatchNorm2d(out_size)
        if attention:
            self.ca = ChannelAttention(out_size)
            self.sa = SpatialAttention()

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        if self.batch_norm:
            out = self.bn1(out)
        out = self.relu(out)
        if self.dropout:
            out = self.drop(out)

        out = self.conv2(out)
        if self.batch_norm:
            out = self.bn2(out)
        if self.attention:
            out = self.ca(out) * out
            out = self.sa(out) * out

        if self.dropout:
            out = self.drop(out)
        if self.last_relu:
            out = self.relu(out)

        return out

class CNNEncoder(nn.Module):
    def __init__(self, in_channels=1, feature_channels=[64,64,64],
                 padding=True, batch_norm=True, max_pool=True,
                 lr_equalization=False, dropout=False, attention=False,
                 global_avg_pool=True):
        super(CNNEncoder,self).__init__()
        print(in_channels)

        self.in_channels = in_channels

        self.depth = len(feature_channels)
        self.max_pool = max_pool
        self.feature_channels = feature_channels
        self.global_avg_pool = global_avg_pool

        assert(len(feature_channels)>0,
               'Error CNNEncoder must have at least 1 conv layer')

        print(f"feature_channels={feature_channels}")

        self.net = []
        prev_channels = in_channels

        for channels in self.feature_channels:
            self.net.append(
                DoubleConvBlock(prev_channels, channels,
                                padding, batch_norm,
                                lr_equalization, dropout, attention)
            )
            if self.max_pool:
                self.net.append(
                    nn.AvgPool2d(kernel_size=2)
                )
            prev_channels = channels

        self.net = nn.Sequential(*self.net)

    def get_out_shape(self,h,w):
        return self.forward(torch.rand(1,self.in_channels,h,w)).shape[1:]

    def forward(self, input):
        out = self.net(input)

        if self.global_avg_pool:
            out = torch.nn.AvgPool2d((out.shape[2], out.shape[3]))(out)
        return out

    
class CNNDecoder(nn.Module):
    def __init__(self, in_channels=1, feature_channels=[64,64,64],
                 padding=True, batch_norm=True, max_pool=True,
                 lr_equalization=False, dropout=False, attention=False,
                 global_avg_pool=True):
        super(CNNDecoder,self).__init__()
        print(in_channels)

        self.in_channels = in_channels

        self.depth = len(feature_channels)
        self.max_pool = max_pool
        self.feature_channels = feature_channels
        self.global_avg_pool = global_avg_pool

        assert(len(feature_channels)>0,
               'Error CNNDecoder must have at least 1 deconv layer')

        print(f"feature_channels={feature_channels}")

        self.net = []
        prev_channels = in_channels

        for i, channels in enumerate(self.feature_channels):
            self.net.append(
                DoubleDeconvBlock(prev_channels, channels,
                                padding, batch_norm,
                                lr_equalization, dropout, attention, last_relu= i<(len(self.feature_channels)-1))
            )
            if self.max_pool:
                self.net.append(
                    nn.Upsample(scale_factor=2, mode='bilinear')
                )
            prev_channels = channels

        self.net = nn.Sequential(*self.net)

    def get_out_shape(self,h,w):
        return self.forward(torch.rand(1,self.in_channels,h,w)).shape[1:]

    def forward(self, input):
        out = self.net(input)

        if self.global_avg_pool:
            out = nn.Upsample(scale_factor=2, mode='bilinear')(out)
        return out
    
    
    
from torch import nn
class UNET(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()

        self.conv1 = self.contract_block(in_channels, 32, 3, 1)
        self.conv2 = self.contract_block(32, 64, 3, 1)
        self.conv3 = self.contract_block(64, 128, 3, 1)

        self.upconv3 = self.expand_block(128, 64, 3, 1)
        self.upconv2 = self.expand_block(64*2, 32, 3, 1)
        self.upconv1 = self.expand_block(32*2, 32, 3, 1)
        
        self.final = self.same_block(32, 1, 3, 1)
        self.res_connection= self.same_block(1, 1, 3, 1)

    def __call__(self, x):

        # downsampling part
        conv1 = self.conv1(x)
        conv2 = self.conv2(conv1)
        conv3 = self.conv3(conv2)

        upconv3 = self.upconv3(conv3)

        upconv2 = self.upconv2(torch.cat([upconv3, conv2], 1))
        upconv1 = self.upconv1(torch.cat([upconv2, conv1], 1))
        output= self.final(upconv1)
        return output#+self.res_connection(x)

    def contract_block(self, in_channels, out_channels, kernel_size, padding):

        contract = nn.Sequential(
            torch.nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding),
            torch.nn.BatchNorm2d(out_channels),
            torch.nn.ReLU(),
            torch.nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding),
            torch.nn.BatchNorm2d(out_channels),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
                                 )

        return contract

    def expand_block(self, in_channels, out_channels, kernel_size, padding):

        expand = nn.Sequential(torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=padding),
                            torch.nn.BatchNorm2d(out_channels),
                            torch.nn.ReLU(),
       
                            torch.nn.Conv2d(out_channels, out_channels, kernel_size, stride=1, padding=padding),
                            torch.nn.BatchNorm2d(out_channels),
                            torch.nn.ReLU(),
                               
                            torch.nn.Upsample( scale_factor=2, mode='bilinear', align_corners=None),
                              
        
                              )
        return expand
               
    def same_block(self, in_channels, out_channels, kernel_size, padding):
            
        same_block=nn.Sequential(torch.nn.Conv2d(in_channels, in_channels, kernel_size, stride=1, padding=padding),
                            torch.nn.BatchNorm2d(in_channels),
                            torch.nn.LeakyReLU(),
       
                            torch.nn.Conv2d(in_channels, in_channels, kernel_size, stride=1, padding=padding),
                            torch.nn.BatchNorm2d(in_channels),
                            torch.nn.LeakyReLU(),
                               
                            
                              
                            torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=padding)
                                )
                 
        return same_block