Upload WNetMSS3D

WNetConfigs.py

@@ -0,0 +1,41 @@
+from transformers import PretrainedConfig
+from typing import List
+
+class WNet3DConfig(PretrainedConfig):
+    model_type = "WNet"
+    def __init__(
+            self,
+            in_ch=1, 
+            out_ch=5, 
+            init_features=64,
+            **kwargs):
+        self.in_ch = in_ch
+        self.out_ch = out_ch
+        self.init_features = init_features
+        super().__init__(**kwargs)
+
+class AttWNet3DConfig(PretrainedConfig):
+    model_type = "AttWNet"
+    def __init__(
+            self,
+            in_ch=1, 
+            out_ch=5, 
+            init_features=64,
+            **kwargs):
+        self.in_ch = in_ch
+        self.out_ch = out_ch
+        self.init_features = init_features
+        super().__init__(**kwargs)
+
+class WNetMSS3DConfig(PretrainedConfig):
+    model_type = "WNetMSS"
+    def __init__(
+            self,
+            in_ch=1, 
+            out_ch=5, 
+            init_features=64,
+            **kwargs):
+        self.in_ch = in_ch
+        self.out_ch = out_ch
+        self.init_features = init_features
+        super().__init__(**kwargs)

WNets.py

@@ -0,0 +1,36 @@
+from transformers import PreTrainedModel
+from .w_net_3d import WNet3dUNet, WNet3dAttUNet, WNet3dUNetMSS
+from .WNetConfigs import WNet3DConfig, AttWNet3DConfig, WNetMSS3DConfig 
+
+class WNet3D(PreTrainedModel):
+    config_class = WNet3DConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = WNet3dUNet(
+            in_ch=config.in_ch, 
+            out_ch=config.out_ch, 
+            init_features=config.init_features)
+    def forward(self, x):
+        return self.model(x)
+    
+class AttWNet3D(PreTrainedModel):
+    config_class = AttWNet3DConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = WNet3dAttUNet(
+            in_ch=config.in_ch, 
+            out_ch=config.out_ch, 
+            init_features=config.init_features)
+    def forward(self, x):
+        return self.model(x)
+    
+class WNetMSS3D(PreTrainedModel):
+    config_class = WNetMSS3DConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = WNet3dUNetMSS(
+            in_ch=config.in_ch, 
+            out_ch=config.out_ch, 
+            init_features=config.init_features)
+    def forward(self, x):
+        return self.model(x)

attention_unet3d.py

@@ -0,0 +1,211 @@
+#!/usr/bin/env python
+
+# from __future__ import print_function, division
+"""
+
+Purpose :
+
+"""
+import torch.nn
+import torch
+import torch.nn as nn
+
+__author__ = "Chethan Radhakrishna and Soumick Chatterjee"
+__credits__ = ["Chethan Radhakrishna", "Soumick Chatterjee"]
+__license__ = "GPL"
+__version__ = "1.0.0"
+__maintainer__ = "Chethan Radhakrishna"
+__email__ = "[email protected]"
+__status__ = "Development"
+
+
+class ConvBlock(nn.Module):
+    """
+    Convolution Block
+    """
+
+    def __init__(self, in_channels, out_channels, k_size=3, stride=1, padding=1, bias=True):
+        super(ConvBlock, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=k_size,
+                      stride=stride, padding=padding, bias=bias),
+            nn.PReLU(num_parameters=out_channels, init=0.25),
+            # nn.Dropout3d(),
+            nn.BatchNorm3d(num_features=out_channels),
+            nn.Conv3d(in_channels=out_channels, out_channels=out_channels, kernel_size=k_size,
+                      stride=stride, padding=padding, bias=bias),
+            nn.PReLU(num_parameters=out_channels, init=0.25),
+            # nn.Dropout3d(),
+            nn.BatchNorm3d(num_features=out_channels))
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+
+class SeparableConvBlock(nn.Module):
+    """
+    Convolution Block
+    """
+
+    def __init__(self, in_channels, out_channels, k_size=3, stride=1, padding=1, bias=True):
+        super(SeparableConvBlock, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=1,
+                      bias=bias),
+            nn.Conv3d(in_channels=out_channels, out_channels=out_channels, kernel_size=k_size,
+                      stride=stride, padding=padding, bias=bias),
+            nn.PReLU(num_parameters=out_channels, init=0.25),
+            # nn.Dropout3d(),
+            nn.BatchNorm3d(num_features=out_channels),
+            nn.Conv3d(in_channels=out_channels, out_channels=out_channels, kernel_size=1,
+                      bias=bias),
+            nn.Conv3d(in_channels=out_channels, out_channels=out_channels, kernel_size=k_size,
+                      stride=stride, padding=padding, bias=bias),
+            nn.PReLU(num_parameters=out_channels, init=0.25),
+            # nn.Dropout3d(),
+            nn.BatchNorm3d(num_features=out_channels))
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+
+class UpConv(nn.Module):
+    """
+    Up Convolution Block
+    """
+
+    # def __init__(self, in_ch, out_ch):
+    def __init__(self, in_channels, out_channels, k_size=3, stride=1, padding=1):
+        super(UpConv, self).__init__()
+        self.up = nn.Sequential(
+            nn.Upsample(scale_factor=2),
+            nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=k_size,
+                      stride=stride, padding=padding),
+            nn.BatchNorm3d(num_features=out_channels),
+            nn.PReLU(num_parameters=out_channels, init=0.25))
+
+    def forward(self, x):
+        x = self.up(x)
+        return x
+
+
+class AttentionBlock(nn.Module):
+    """
+    Attention Block
+    """
+
+    def __init__(self, f_g, f_l, f_int):
+        super(AttentionBlock, self).__init__()
+
+        self.W_g = nn.Sequential(
+            nn.Conv3d(f_l, f_int, kernel_size=1, stride=1, padding=0, bias=True),
+            nn.BatchNorm3d(f_int)
+        )
+
+        self.W_x = nn.Sequential(
+            nn.Conv3d(f_g, f_int, kernel_size=1, stride=1, padding=0, bias=True),
+            nn.BatchNorm3d(f_int)
+        )
+
+        self.psi = nn.Sequential(
+            nn.Conv3d(f_int, 1, kernel_size=1, stride=1, padding=0, bias=True),
+            nn.BatchNorm3d(1),
+            nn.Sigmoid()
+        )
+
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, g, x):
+        g1 = self.W_g(g)
+        x1 = self.W_x(x)
+        psi = self.relu(g1 + x1)
+        psi = self.psi(psi)
+        out = x * psi
+        return out
+
+
+class AttUnet(nn.Module):
+    """
+    Attention Unet implementation
+    Paper: https://arxiv.org/abs/1804.03999
+    """
+
+    def __init__(self, in_ch=1, out_ch=6, init_features=64):
+        super(AttUnet, self).__init__()
+
+        n1 = init_features
+        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]
+
+        self.Maxpool1 = nn.MaxPool3d(kernel_size=2, stride=2)
+        self.Maxpool2 = nn.MaxPool3d(kernel_size=2, stride=2)
+        self.Maxpool3 = nn.MaxPool3d(kernel_size=2, stride=2)
+        self.Maxpool4 = nn.MaxPool3d(kernel_size=2, stride=2)
+
+        self.Conv1 = ConvBlock(in_ch, filters[0])
+        self.Conv2 = SeparableConvBlock(filters[0], filters[1])
+        self.Conv3 = SeparableConvBlock(filters[1], filters[2])
+        self.Conv4 = SeparableConvBlock(filters[2], filters[3])
+        self.Conv5 = SeparableConvBlock(filters[3], filters[4])
+
+        self.Up5 = UpConv(filters[4], filters[3])
+        self.Att5 = AttentionBlock(f_g=filters[3], f_l=filters[3], f_int=filters[2])
+        self.Up_conv5 = SeparableConvBlock(filters[4], filters[3])
+
+        self.Up4 = UpConv(filters[3], filters[2])
+        self.Att4 = AttentionBlock(f_g=filters[2], f_l=filters[2], f_int=filters[1])
+        self.Up_conv4 = SeparableConvBlock(filters[3], filters[2])
+
+        self.Up3 = UpConv(filters[2], filters[1])
+        self.Att3 = AttentionBlock(f_g=filters[1], f_l=filters[1], f_int=filters[0])
+        self.Up_conv3 = SeparableConvBlock(filters[2], filters[1])
+
+        self.Up2 = UpConv(filters[1], filters[0])
+        self.Att2 = AttentionBlock(f_g=filters[0], f_l=filters[0], f_int=32)
+        self.Up_conv2 = ConvBlock(filters[1], filters[0])
+
+        self.Conv = nn.Conv3d(filters[0], out_ch, kernel_size=1, stride=1, padding=0)
+
+        # self.active = torch.nn.Sigmoid()
+
+    def forward(self, x):
+        e1 = self.Conv1(x)
+
+        e2 = self.Maxpool1(e1)
+        e2 = self.Conv2(e2)
+
+        e3 = self.Maxpool2(e2)
+        e3 = self.Conv3(e3)
+
+        e4 = self.Maxpool3(e3)
+        e4 = self.Conv4(e4)
+
+        e5 = self.Maxpool4(e4)
+        e5 = self.Conv5(e5)
+
+        d5 = self.Up5(e5)
+        x4 = self.Att5(d5, e4)
+        d5 = torch.cat((x4, d5), dim=1)
+        d5 = self.Up_conv5(d5)
+
+        d4 = self.Up4(d5)
+        x3 = self.Att4(d4, e3)
+        d4 = torch.cat((x3, d4), dim=1)
+        d4 = self.Up_conv4(d4)
+
+        d3 = self.Up3(d4)
+        x2 = self.Att3(d3, e2)
+        d3 = torch.cat((x2, d3), dim=1)
+        d3 = self.Up_conv3(d3)
+
+        d2 = self.Up2(d3)
+        x1 = self.Att2(d2, e1)
+        d2 = torch.cat((x1, d2), dim=1)
+        d2 = self.Up_conv2(d2)
+
+        out = self.Conv(d2)
+
+        #  out = self.active(out)
+
+        return out

config.json

@@ -0,0 +1,15 @@
+{
+  "architectures": [
+    "WNetMSS3D"
+  ],
+  "auto_map": {
+    "AutoConfig": "WNetConfigs.WNetMSS3DConfig",
+    "AutoModel": "WNets.WNetMSS3D"
+  },
+  "in_ch": 1,
+  "init_features": 64,
+  "model_type": "WNetMSS",
+  "out_ch": 5,
+  "torch_dtype": "float32",
+  "transformers_version": "4.44.2"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7116678da6ee0dfe953f794ce2b65e69e17693d874e76007c9e0489267fda23f
+size 929631232

unet3d.py

@@ -0,0 +1,321 @@
+#!/usr/bin/env python
+
+# from __future__ import print_function, division
+"""
+
+Purpose : 
+
+"""
+
+import torch
+import torch.nn as nn
+import torch.utils.data
+
+__author__ = "Kartik Prabhu, Mahantesh Pattadkal, and Soumick Chatterjee"
+__copyright__ = "Copyright 2020, Faculty of Computer Science, Otto von Guericke University Magdeburg, Germany"
+__credits__ = ["Kartik Prabhu", "Mahantesh Pattadkal", "Soumick Chatterjee"]
+__license__ = "GPL"
+__version__ = "1.0.0"
+__maintainer__ = "Soumick Chatterjee"
+__email__ = "[email protected]"
+__status__ = "Production"
+
+
+class ConvBlock(nn.Module):
+    """
+    Convolution Block
+    """
+
+    def __init__(self, in_channels, out_channels, k_size=3, stride=1, padding=1, bias=True):
+        super(ConvBlock, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=k_size,
+                      stride=stride, padding=padding, bias=bias),
+            nn.BatchNorm3d(num_features=out_channels),
+            nn.LeakyReLU(inplace=True),
+            nn.Conv3d(in_channels=out_channels, out_channels=out_channels, kernel_size=k_size,
+                      stride=stride, padding=padding, bias=bias),
+            nn.BatchNorm3d(num_features=out_channels),
+            nn.LeakyReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+
+class SeparableConvBlock(nn.Module):
+    """
+    Convolution Block
+    """
+
+    def __init__(self, in_channels, out_channels, k_size=3, stride=1, padding=1, bias=True):
+        super(SeparableConvBlock, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=1,
+                      bias=bias),
+            nn.Conv3d(in_channels=out_channels, out_channels=out_channels, kernel_size=k_size,
+                      stride=stride, padding=padding, bias=bias),
+            nn.PReLU(num_parameters=out_channels, init=0.25),
+            nn.Dropout3d(),
+            nn.BatchNorm3d(num_features=out_channels),
+            nn.Conv3d(in_channels=out_channels, out_channels=out_channels, kernel_size=1,
+                      bias=bias),
+            nn.Conv3d(in_channels=out_channels, out_channels=out_channels, kernel_size=k_size,
+                      stride=stride, padding=padding, bias=bias),
+            nn.PReLU(num_parameters=out_channels, init=0.25),
+            nn.Dropout3d(),
+            nn.BatchNorm3d(num_features=out_channels))
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+
+class UpConv(nn.Module):
+    """
+    Up Convolution Block
+    """
+
+    # def __init__(self, in_ch, out_ch):
+    def __init__(self, in_channels, out_channels, k_size=3, stride=1, padding=1, bias=True):
+        super(UpConv, self).__init__()
+        self.up = nn.Sequential(
+            nn.Upsample(scale_factor=2),
+            nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=k_size,
+                      stride=stride, padding=padding, bias=bias),
+            nn.BatchNorm3d(num_features=out_channels),
+            nn.LeakyReLU(inplace=True))
+
+    def forward(self, x):
+        x = self.up(x)
+        return x
+
+
+class UNet(nn.Module):
+    """
+    UNet - Basic Implementation
+    Input _ [batch * channel(# of channels of each image) * depth(# of frames) * height * width].
+    Paper : https://arxiv.org/abs/1505.04597
+    """
+
+    def __init__(self, in_ch=1, out_ch=1, init_features=64):
+        super(UNet, self).__init__()
+
+        n1 = init_features  
+        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]  # 64,128,256,512,1024
+
+        self.Maxpool1 = nn.MaxPool3d(kernel_size=2, stride=2)
+        self.Maxpool2 = nn.MaxPool3d(kernel_size=2, stride=2)
+        self.Maxpool3 = nn.MaxPool3d(kernel_size=2, stride=2)
+        self.Maxpool4 = nn.MaxPool3d(kernel_size=2, stride=2)
+
+        self.Conv1 = ConvBlock(in_ch, filters[0])
+        self.Conv2 = SeparableConvBlock(filters[0], filters[1])
+        self.Conv3 = SeparableConvBlock(filters[1], filters[2])
+        self.Conv4 = SeparableConvBlock(filters[2], filters[3])
+        self.Conv5 = SeparableConvBlock(filters[3], filters[4])
+
+        self.Up5 = UpConv(filters[4], filters[3])
+        self.Up_conv5 = SeparableConvBlock(filters[4], filters[3])
+
+        self.Up4 = UpConv(filters[3], filters[2])
+        self.Up_conv4 = SeparableConvBlock(filters[3], filters[2])
+
+        self.Up3 = UpConv(filters[2], filters[1])
+        self.Up_conv3 = SeparableConvBlock(filters[2], filters[1])
+
+        self.Up2 = UpConv(filters[1], filters[0])
+        self.Up_conv2 = ConvBlock(filters[1], filters[0])
+
+        self.Conv = nn.Conv3d(filters[0], out_ch, kernel_size=1, stride=1, padding=0)
+
+    # self.active = torch.nn.Sigmoid()
+
+    def forward(self, x):
+        # print("unet")
+        # print(x.shape)
+        # print(padded.shape)
+
+        e1 = self.Conv1(x)
+        # print("conv1:")
+        # print(e1.shape)
+
+        e2 = self.Maxpool1(e1)
+        e2 = self.Conv2(e2)
+        # print("conv2:")
+        # print(e2.shape)
+
+        e3 = self.Maxpool2(e2)
+        e3 = self.Conv3(e3)
+        # print("conv3:")
+        # print(e3.shape)
+
+        e4 = self.Maxpool3(e3)
+        e4 = self.Conv4(e4)
+        # print("conv4:")
+        # print(e4.shape)
+
+        e5 = self.Maxpool4(e4)
+        e5 = self.Conv5(e5)
+        # print("conv5:")
+        # print(e5.shape)
+
+        d5 = self.Up5(e5)
+        # print("d5:")
+        # print(d5.shape)
+        # print("e4:")
+        # print(e4.shape)
+        d5 = torch.cat((e4, d5), dim=1)
+        d5 = self.Up_conv5(d5)
+        # print("upconv5:")
+        # print(d5.size)
+
+        d4 = self.Up4(d5)
+        # print("d4:")
+        # print(d4.shape)
+        d4 = torch.cat((e3, d4), dim=1)
+        d4 = self.Up_conv4(d4)
+        # print("upconv4:")
+        # print(d4.shape)
+        d3 = self.Up3(d4)
+        d3 = torch.cat((e2, d3), dim=1)
+        d3 = self.Up_conv3(d3)
+        # print("upconv3:")
+        # print(d3.shape)
+        d2 = self.Up2(d3)
+        d2 = torch.cat((e1, d2), dim=1)
+        d2 = self.Up_conv2(d2)
+        # print("upconv2:")
+        # print(d2.shape)
+        out = self.Conv(d2)
+        # print("out:")
+        # print(out.shape)
+        # d1 = self.active(out)
+
+        return out
+
+
+class UNetDeepSup(nn.Module):
+    """
+    UNet - Basic Implementation
+    Input _ [batch * channel(# of channels of each image) * depth(# of frames) * height * width].
+    Paper : https://arxiv.org/abs/1505.04597
+    """
+
+    def __init__(self, in_ch=1, out_ch=1, init_features=64):
+        super(UNetDeepSup, self).__init__()
+
+        n1 = init_features
+        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]  # 64,128,256,512,1024
+
+        self.Maxpool1 = nn.MaxPool3d(kernel_size=2, stride=2)
+        self.Maxpool2 = nn.MaxPool3d(kernel_size=2, stride=2)
+        self.Maxpool3 = nn.MaxPool3d(kernel_size=2, stride=2)
+        self.Maxpool4 = nn.MaxPool3d(kernel_size=2, stride=2)
+
+        self.Conv1 = ConvBlock(in_ch, filters[0])
+        self.Conv2 = SeparableConvBlock(filters[0], filters[1])
+        self.Conv3 = SeparableConvBlock(filters[1], filters[2])
+        self.Conv4 = SeparableConvBlock(filters[2], filters[3])
+        self.Conv5 = SeparableConvBlock(filters[3], filters[4])
+
+        # 1x1x1 Convolution for Deep Supervision
+        self.Conv_d3 = SeparableConvBlock(filters[1], 1)
+        self.Conv_d4 = SeparableConvBlock(filters[2], 1)
+
+        self.Up5 = UpConv(filters[4], filters[3])
+        self.Up_conv5 = SeparableConvBlock(filters[4], filters[3])
+
+        self.Up4 = UpConv(filters[3], filters[2])
+        self.Up_conv4 = SeparableConvBlock(filters[3], filters[2])
+
+        self.Up3 = UpConv(filters[2], filters[1])
+        self.Up_conv3 = SeparableConvBlock(filters[2], filters[1])
+
+        self.Up2 = UpConv(filters[1], filters[0])
+        self.Up_conv2 = ConvBlock(filters[1], filters[0])
+
+        self.Conv = nn.Conv3d(filters[0], out_ch, kernel_size=1, stride=1, padding=0)
+
+        for submodule in self.modules():
+            submodule.register_forward_hook(self.nan_hook)
+
+    # self.active = torch.nn.Sigmoid()
+
+    def nan_hook(self, module, inp, output):
+        for i, out in enumerate(output):
+            nan_mask = torch.isnan(out)
+            if nan_mask.any():
+                print("In", self.__class__.__name__)
+                torch.save(inp, '/nfs1/sutrave/outputs/nan_values_input/inp_2_Nov.pt')
+                raise RuntimeError(" classname " + self.__class__.__name__ + "i " + str(
+                    i) + f" module: {module} classname {self.__class__.__name__} Found NAN in output {i} at indices: ",
+                                   nan_mask.nonzero(), "where:", out[nan_mask.nonzero()[:, 0].unique(sorted=True)])
+
+    def forward(self, x):
+        # print("unet")
+        # print(x.shape)
+        # print(padded.shape)
+
+        e1 = self.Conv1(x)
+        # print("conv1:")
+        # print(e1.shape)
+
+        e2 = self.Maxpool1(e1)
+        e2 = self.Conv2(e2)
+        # print("conv2:")
+        # print(e2.shape)
+
+        e3 = self.Maxpool2(e2)
+        e3 = self.Conv3(e3)
+        # print("conv3:")
+        # print(e3.shape)
+
+        e4 = self.Maxpool3(e3)
+        e4 = self.Conv4(e4)
+        # print("conv4:")
+        # print(e4.shape)
+
+        e5 = self.Maxpool4(e4)
+        e5 = self.Conv5(e5)
+        # print("conv5:")
+        # print(e5.shape)
+
+        d5 = self.Up5(e5)
+        # print("d5:")
+        # print(d5.shape)
+        # print("e4:")
+        # print(e4.shape)
+        d5 = torch.cat((e4, d5), dim=1)
+        d5 = self.Up_conv5(d5)
+        # print("upconv5:")
+        # print(d5.size)
+
+        d4 = self.Up4(d5)
+        # print("d4:")
+        # print(d4.shape)
+        d4 = torch.cat((e3, d4), dim=1)
+        d4 = self.Up_conv4(d4)
+        d4_out = self.Conv_d4(d4)
+
+        # print("upconv4:")
+        # print(d4.shape)
+        d3 = self.Up3(d4)
+        d3 = torch.cat((e2, d3), dim=1)
+        d3 = self.Up_conv3(d3)
+        d3_out = self.Conv_d3(d3)
+
+        # print("upconv3:")
+        # print(d3.shape)
+        d2 = self.Up2(d3)
+        d2 = torch.cat((e1, d2), dim=1)
+        d2 = self.Up_conv2(d2)
+        # print("upconv2:")
+        # print(d2.shape)
+        out = self.Conv(d2)
+        # print("out:")
+        # print(out.shape)
+        # d1 = self.active(out)
+
+        return out

w_net_3d.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python
+
+# from __future__ import print_function, division
+"""
+
+Purpose :
+
+"""
+import torch.nn
+import torch
+import torch.nn as nn
+from .attention_unet3d import AttUnet
+from .unet3d import UNet, UNetDeepSup
+
+__author__ = "Chethan Radhakrishna and Soumick Chatterjee"
+__credits__ = ["Chethan Radhakrishna", "Soumick Chatterjee"]
+__license__ = "GPL"
+__version__ = "1.0.0"
+__maintainer__ = "Chethan Radhakrishna"
+__email__ = "[email protected]"
+__status__ = "Development"
+
+
+class WNet3dAttUNet(nn.Module):
+    """
+    Attention Unet implementation
+    Paper: https://arxiv.org/abs/1804.03999
+    """
+
+    def __init__(self, in_ch=1, out_ch=6, init_features=64):
+        super(WNet3dAttUNet, self).__init__()
+
+        self.Encoder = AttUnet(in_ch=in_ch, out_ch=out_ch, init_features=init_features)
+        self.Decoder = AttUnet(in_ch=out_ch, out_ch=in_ch, init_features=init_features)
+
+        self.activation = torch.nn.Softmax(dim=1)
+
+        self.Conv = nn.Conv3d(out_ch, in_ch, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, ip, ip_mask=None, ops="both"):
+        encoder_op = self.Encoder(ip)
+        if ip_mask is not None:
+            encoder_op = ip_mask * encoder_op
+        class_prob = self.activation(encoder_op)
+        feature_rep = self.Conv(encoder_op)
+        if ops == "enc":
+            return class_prob, feature_rep
+        reconstructed_op = self.Decoder(class_prob)
+        # if ip_mask is not None:
+        #     reconstructed_op = torch.amax(ip_mask, dim=1, keepdim=True) * reconstructed_op
+        if ops == "dec":
+            return reconstructed_op
+        if ops == "both":
+            return class_prob, feature_rep, reconstructed_op
+        else:
+            raise ValueError('Invalid ops, ops must be in [enc, dec, both]')
+
+
+class WNet3dUNet(nn.Module):
+    """
+    Attention Unet implementation
+    Paper: https://arxiv.org/abs/1804.03999
+    """
+
+    def __init__(self, in_ch=1, out_ch=6, init_features=64):
+        super(WNet3dUNet, self).__init__()
+
+        self.Encoder = UNet(in_ch=in_ch, out_ch=out_ch, init_features=init_features)
+        self.Decoder = UNet(in_ch=out_ch, out_ch=in_ch, init_features=init_features)
+
+        self.activation = torch.nn.Softmax(dim=1)
+
+        self.Conv = nn.Conv3d(out_ch, in_ch, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, ip, ip_mask=None, ops="both"):
+        encoder_op = self.Encoder(ip)
+        if ip_mask is not None:
+            encoder_op = ip_mask * encoder_op
+        class_prob = self.activation(encoder_op)
+        feature_rep = self.Conv(encoder_op)
+        if ops == "enc":
+            return class_prob, feature_rep
+        reconstructed_op = self.Decoder(class_prob)
+        # if ip_mask is not None:
+        #     reconstructed_op = torch.amax(ip_mask, dim=1, keepdim=True) * reconstructed_op
+        if ops == "dec":
+            return reconstructed_op
+        if ops == "both":
+            return class_prob, feature_rep, reconstructed_op
+        else:
+            raise ValueError('Invalid ops, ops must be in [enc, dec, both]')
+
+
+class WNet3dUNetMSS(nn.Module):
+    """
+    Attention Unet implementation
+    Paper: https://arxiv.org/abs/1804.03999
+    """
+
+    def __init__(self, in_ch=1, out_ch=6, init_features=64):
+        super(WNet3dUNetMSS, self).__init__()
+
+        self.Encoder = UNetDeepSup(in_ch=in_ch, out_ch=out_ch, init_features=init_features)
+        self.Decoder = UNetDeepSup(in_ch=out_ch, out_ch=in_ch, init_features=init_features)
+
+        self.activation = torch.nn.Softmax(dim=1)
+
+        self.Conv = nn.Conv3d(out_ch, in_ch, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, ip, ip_mask=None, ops="both"):
+        encoder_op = self.Encoder(ip)
+        if ip_mask is not None:
+            encoder_op = ip_mask * encoder_op
+        class_prob = self.activation(encoder_op)
+        feature_rep = self.Conv(encoder_op)
+        if ops == "enc":
+            return class_prob, feature_rep
+        reconstructed_op = self.Decoder(class_prob)
+        # if ip_mask is not None:
+        #     reconstructed_op = torch.amax(ip_mask, dim=1, keepdim=True) * reconstructed_op
+        if ops == "dec":
+            return reconstructed_op
+        if ops == "both":
+            return class_prob, feature_rep, reconstructed_op
+        else:
+            raise ValueError('Invalid ops, ops must be in [enc, dec, both]')