From 099432af033f8cf18325f076f646f69f4fd9d549 Mon Sep 17 00:00:00 2001 From: Sashi Date: Thu, 3 Jan 2019 00:10:30 -0600 Subject: [PATCH] + added anatomy net for segmentation + formatting change --- core/NeuralArchitectures/__init__.py | 6 +- core/NeuralArchitectures/contextnet.py | 147 ++++++++++--- core/NeuralArchitectures/pointnet.py | 64 ++++-- core/NeuralArchitectures/unet.py | 294 ++++++++++++++++--------- core/NeuralLayers/__init__.py | 3 +- core/NeuralLayers/lossfunctions.py | 68 ++++-- 6 files changed, 397 insertions(+), 185 deletions(-) diff --git a/core/NeuralArchitectures/__init__.py b/core/NeuralArchitectures/__init__.py index fbc0fb9..4ee08e1 100644 --- a/core/NeuralArchitectures/__init__.py +++ b/core/NeuralArchitectures/__init__.py @@ -5,7 +5,7 @@ "ResidualNet", "InceptionV4", "DenseNet", "LinearVAE", "ConvolutionalVAE", "PGGAN", "ContextNet", "FeatureNet", "FeatureCapNet", - "PointNet", "UNet", "UNetPatch", "UNetMini", + "PointNet", "UNet", "ANet", "NeuralDecisionForest", "Models"] from .capsulenet import CapsuleNet @@ -20,15 +20,13 @@ from .convolutionalvae import ConvolutionalVAE from .pggan import PGGAN from .contextnet import ContextNet -from .featurenet import FeatureNet, FeatureCapNet from .pointnet import PointNet -from .unet import UNet, UNetMini +from .unet import UNet, ANet from .trees import NeuralDecisionForest del trees del unet del pointnet -del featurenet del contextnet del pggan del capsulenet diff --git a/core/NeuralArchitectures/contextnet.py b/core/NeuralArchitectures/contextnet.py index e58141c..4aa3310 100644 --- a/core/NeuralArchitectures/contextnet.py +++ b/core/NeuralArchitectures/contextnet.py @@ -1,51 +1,126 @@ - -""" TensorMONK's :: NeuralArchitectures """ - -import torch +""" TensorMONK's :: NeuralArchitectures """ import torch.nn as nn -import numpy as np -from ..NeuralLayers import * -#==============================================================================# +from ..NeuralLayers import Convolution +from ..NeuralLayers import ContextNet_Bottleneck class ContextNet(nn.Module): - """ - Implemented https://arxiv.org/pdf/1805.04554.pdf + r"""ContextNet: Exploring Context and Detail for Semantic Seg in Real-time + Implemented from https://arxiv.org/pdf/1805.04554.pdf + + Naming Convention for layers in the module: + cnv::Convolution:: regular convolution block + bn::bottleneck:: bottlenect residual block + dw::Convilution:: depth-wise seperable convolution block + dn:: deep netwrk for Context + sn:: shallow network for Context + Args: + tensor_size: shape of tensor in BCHW """ def __init__(self, tensor_size=(1, 3, 1024, 2048), *args, **kwargs): super(ContextNet, self).__init__() - Strides = [2, 1, 1] - bottleneck = ContextNet_Bottleneck - normalization = "batch" - self.DeepNET = nn.Sequential() + normalization, strides = "batch", [2, 1, 1] + bottleneck = ContextNet_Bottleneck + + self.DeepNET = nn.Sequential() self.ShallowNET = nn.Sequential() - self.DeepNET.add_module("AVGPL", nn.AvgPool2d((5,5), (4,4), 2)) # 1, 1, 256, 512 - self.DeepNET.add_module("DN_CNV1", Convolution(tensor_size, 3, 32, 2, True, "relu", 0., normalization, False, 1)) # 1, 1, 128, 256 - self.DeepNET.add_module("DN_BN10", bottleneck(self.DeepNET[-1].tensor_size, 3, 32, 1, expansion=1)) - self.DeepNET.add_module("DN_BN20", bottleneck(self.DeepNET[-1].tensor_size, 3, 32, 1, expansion=6)) # 1, 1, 128, 256 - for i in range(3): self.DeepNET.add_module("DN_BN3"+str(i), bottleneck(self.DeepNET[-1].tensor_size, 3, 48, Strides[i], expansion=6)) # 1, 1, 64, 128 - for i in range(3): self.DeepNET.add_module("DN_BN4"+str(i), bottleneck(self.DeepNET[-1].tensor_size, 3, 64, Strides[i], expansion=6)) # 1, 1, 32, 64 - for i in range(2): self.DeepNET.add_module("DN_BN5"+str(i), bottleneck(self.DeepNET[-1].tensor_size, 3, 96, 1, expansion=6)) - for i in range(2): self.DeepNET.add_module("DN_BN6"+str(i), bottleneck(self.DeepNET[-1].tensor_size, 3, 128, 1, expansion=6)) # 1, 1, 32, 64 - self.DeepNET.add_module("DN_CNV2", Convolution(self.DeepNET[-1].tensor_size, 3, 128, 1, True, "relu", 0., normalization, False, 1)) # 1, 1, 32, 64 - self.DeepNET.add_module("UPSMPLE", nn.Upsample(scale_factor = 4, mode = 'bilinear')) # 1, 1, 128, 256 - _tensor_size = (1, 128, self.DeepNET[-2].tensor_size[2]*4, self.DeepNET[-2].tensor_size[3]*4) - self.DeepNET.add_module("DN_DW11", Convolution(_tensor_size, 3, _tensor_size[1], 1, True, "relu", 0., None, False, groups =_tensor_size[1], dilation = 4)) - self.DeepNET.add_module("DN_DW12", Convolution(self.DeepNET[-1].tensor_size, 1, 128, 1, True, "relu", 0., normalization, False, 1)) - self.DeepNET.add_module("DN_CNV3", Convolution(self.DeepNET[-1].tensor_size, 1, 128, 1, True, "relu", 0., normalization, False, 1)) # 128, 256 + # 1, 1, 256, 512 + self.DeepNET.add_module("avgpl", nn.AvgPool2d((5, 5), (4, 4), 2)) + # 1, 1, 128, 256 + self.DeepNET.add_module("dn_cnv1", Convolution(tensor_size, 3, 32, 2, + True, "relu", 0., normalization, False, 1)) + self.DeepNET.add_module("dn_bn10", + bottleneck(self.DeepNET[-1].tensor_size, 3, 32, + 1, expansion=1)) + self.DeepNET.add_module("dn_bn20", + bottleneck(self.DeepNET[-1].tensor_size, 3, 32, + 1, expansion=6)) # 1, 1, 128, 256 + for i in range(3): + self.DeepNET.add_module("dn_bn3"+str(i), + bottleneck(self.DeepNET[-1].tensor_size, 3, + 48, strides[i], expansion=6)) + # 1, 1, 64, 128 + for i in range(3): + self.DeepNET.add_module("dn_bn4"+str(i), + bottleneck(self.DeepNET[-1].tensor_size, 3, + 64, strides[i], expansion=6)) + # 1, 1, 32, 64 + for i in range(2): + self.DeepNET.add_module("dn_bn5"+str(i), + bottleneck(self.DeepNET[-1].tensor_size, 3, + 96, 1, expansion=6)) + for i in range(2): + self.DeepNET.add_module("DN_BN6"+str(i), + bottleneck(self.DeepNET[-1].tensor_size, 3, + 128, 1, expansion=6)) + # 1, 1, 32, 64 + self.DeepNET.add_module("dn_cnv2", + Convolution(self.DeepNET[-1].tensor_size, 3, + 128, 1, True, "relu", 0., + normalization, False, 1)) + # 1, 1, 32, 64 + self.DeepNET.add_module("upsample", nn.Upsample(scale_factor=4, + mode='bilinear')) + # 1, 1, 128, 256 + _tensor_size = (1, 128, self.DeepNET[-2].tensor_size[2]*4, + self.DeepNET[-2].tensor_size[3]*4) + self.DeepNET.add_module("dn_dw11", + Convolution(_tensor_size, 3, _tensor_size[1], + 1, True, "relu", 0., None, False, + groups=_tensor_size[1], + dilation=4)) + self.DeepNET.add_module("dn_dw12", + Convolution(self.DeepNET[-1].tensor_size, 1, + 128, 1, True, "relu", 0., + normalization, False, 1)) + self.DeepNET.add_module("dn_cnv3", + Convolution(self.DeepNET[-1].tensor_size, 1, + 128, 1, True, "relu", 0., + normalization, False, 1)) + # 128, 256 activation, pre_nm, groups = "relu", False, 1 - self.ShallowNET.add_module("SM_CNV1", Convolution(tensor_size, 3, 32, 2, True, "relu", 0.,True, False, 1)) # 512 x 1024 - self.ShallowNET.add_module("SM_DW11", Convolution(self.ShallowNET[-1].tensor_size, 3, 32, 2, True, activation, 0., None, pre_nm, groups = tensor_size[1])) # 256, 512 - self.ShallowNET.add_module("SM_DW12", Convolution(self.ShallowNET[-1].tensor_size, 1, 64, 1,True, activation, 0., normalization, pre_nm, groups)) - self.ShallowNET.add_module("SM_DW21", Convolution(self.ShallowNET[-1].tensor_size, 3, 64, 2, True, activation, 0., None, pre_nm, groups = tensor_size[1])) # 128, 256 - self.ShallowNET.add_module("SM_DW22", Convolution(self.ShallowNET[-1].tensor_size, 1, 128, 1,True, activation, 0., normalization, pre_nm, groups)) - self.ShallowNET.add_module("SM_DW31", Convolution(self.ShallowNET[-1].tensor_size, 3, 128, 1, True, activation, 0., None, pre_nm, groups = tensor_size[1])) - self.ShallowNET.add_module("SM_DW32", Convolution(self.ShallowNET[-1].tensor_size, 1, 128, 1,True, activation, 0., normalization, pre_nm, groups)) - self.ShallowNET.add_module("SM_CNV2", Convolution(self.ShallowNET[-1].tensor_size, 1, 128, 1,True, activation, 0., normalization, pre_nm, groups)) # 128, 256 + self.ShallowNET.add_module("sm_cnv1", + Convolution(tensor_size, 3, 32, 2, True, + "relu", 0., True, False, 1)) + # 512 x 1024 + self.ShallowNET.add_module("sm_dw11", + Convolution(self.ShallowNET[-1].tensor_size, + 3, 32, 2, True, activation, 0., + None, pre_nm, + groups=tensor_size[1])) + # 256, 512 + self.ShallowNET.add_module("sm_dw12", + Convolution(self.ShallowNET[-1].tensor_size, + 1, 64, 1, True, activation, 0., + normalization, pre_nm, groups)) + self.ShallowNET.add_module("sm_dw21", + Convolution(self.ShallowNET[-1].tensor_size, + 3, 64, 2, True, activation, 0., + None, pre_nm, + groups=tensor_size[1])) + self.ShallowNET.add_module("sm_dw22", + Convolution(self.ShallowNET[-1].tensor_size, + 1, 128, 1, True, activation, 0., + normalization, pre_nm, groups)) + self.ShallowNET.add_module("sm_dw31", + Convolution(self.ShallowNET[-1].tensor_size, + 3, 128, 1, True, activation, 0., + None, pre_nm, + groups=tensor_size[1])) + self.ShallowNET.add_module("sm_dw32", + Convolution(self.ShallowNET[-1].tensor_size, + 1, 128, 1, True, activation, 0., + normalization, pre_nm, groups)) + self.ShallowNET.add_module("sm_cnv2", + Convolution(self.ShallowNET[-1].tensor_size, + 1, 128, 1, True, activation, 0., + normalization, pre_nm, groups)) + # 128, 256 self.tensor_size = self.ShallowNET[-1].tensor_size - self.FuseNET = Convolution(self.ShallowNET[-1].tensor_size, 1, self.ShallowNET[-1].tensor_size[1], 1, True) + self.FuseNET = Convolution(self.ShallowNET[-1].tensor_size, 1, + self.ShallowNET[-1].tensor_size[1], 1, True) def forward(self, tensor): return self.FuseNET(self.DeepNET(tensor)+self.ShallowNET(tensor)) diff --git a/core/NeuralArchitectures/pointnet.py b/core/NeuralArchitectures/pointnet.py index 740bcb4..a9fffc9 100644 --- a/core/NeuralArchitectures/pointnet.py +++ b/core/NeuralArchitectures/pointnet.py @@ -1,40 +1,58 @@ +""" TensorMONK's :: NeuralArchitectures """ +import torch.nn as nn +from core.NeuralLayers import Convolution +# =========================================================================== # -""" TensorMONK's :: NeuralArchitectures """ -import torch -import torch.nn as nn -import numpy as np -from ..NeuralLayers import * -#==============================================================================# class PointNet(nn.Module): + r""" + Implemented from paper: Learning Discriminative and Transformation + Covariant Local Feature Detectors + Args: + tensor_size: shape of tensor in BCHW + (None/any integer >0, channels, height, width) + out_channels: depth of output feature channels. + activation: None/relu/relu6/lklu/elu/prelu/tanh/sigm/maxo/rmxo/swish + normalization: None/batch/group/instance/layer/pixelwise """ - Implemented http://openaccess.thecvf.com/content_cvpr_2017/papers/Zhang_Learning_Discriminative_and_CVPR_2017_paper.pdf - """ - def __init__(self, tensor_size=(1, 1, 32, 32), out_channels=2, *args, **kwargs): + def __init__(self, tensor_size=(1, 1, 32, 32), out_channels=2, + *args, **kwargs): super(PointNet, self).__init__() normalization = "batch" - self.PointNET = nn.Sequential() - self.PointNET.add_module("CONV1", Convolution(tensor_size, 5, 32, 1, False, "relu", 0., None)) + activation = "relu" + self.PointNET = nn.Sequential() + self.PointNET.add_module("CONV1", + Convolution(tensor_size, 5, 32, 1, False, + activation, 0., None)) self.PointNET.add_module("POOL1", nn.MaxPool2d(2)) _tensor_size = self.PointNET[-2].tensor_size - _tensor_size = (_tensor_size[0], _tensor_size[1], _tensor_size[2]//2, _tensor_size[3]//2) - #print(_tensor_size) - self.PointNET.add_module("CONV2", Convolution(_tensor_size, 3, 128,1, False, "relu", 0.,normalization)) + _tensor_size = (_tensor_size[0], _tensor_size[1], + _tensor_size[2]//2, _tensor_size[3]//2) + self.PointNET.add_module("CONV2", + Convolution(_tensor_size, 3, 128, 1, False, + activation, 0., normalization)) self.PointNET.add_module("POOL2", nn.MaxPool2d(2)) _tensor_size = self.PointNET[-2].tensor_size - _tensor_size = (_tensor_size[0], _tensor_size[1], _tensor_size[2]//2, _tensor_size[3]//2) - #print(_tensor_size) - self.PointNET.add_module("CONV3", Convolution(_tensor_size, 3, 128, 1, False, "relu", 0.,normalization)) - #print(self.PointNET[-1].tensor_size) - self.PointNET.add_module("CONV4", Convolution(self.PointNET[-1].tensor_size, 3, 256, 1, False, "relu", 0.,normalization)) - #print(self.PointNET[-1].tensor_size) - self.PointNET.add_module("CONV5", Convolution(self.PointNET[-1].tensor_size, 2, out_channels, 1, False, "relu", 0.,None)) - print(self.PointNET[-1].tensor_size) + _tensor_size = (_tensor_size[0], _tensor_size[1], + _tensor_size[2]//2, _tensor_size[3]//2) + self.PointNET.add_module("CONV3", + Convolution(_tensor_size, 3, 128, 1, False, + activation, 0., normalization)) + self.PointNET.add_module("CONV4", + Convolution(self.PointNET[-1].tensor_size, 3, + 256, 1, False, activation, 0., + normalization)) + self.PointNET.add_module("CONV5", + Convolution(self.PointNET[-1].tensor_size, 2, + out_channels, 1, False, + activation, 0., None)) self.tensor_size = self.PointNET[-1].tensor_size + def forward(self, tensor): return self.PointNET(tensor).squeeze(2).squeeze(2) -# from core.NeuralLayers import * + +# import torch # tensor_size = (1, 1, 32,32) # test = PointNet(tensor_size) # test(torch.rand(1,1,32,32)) diff --git a/core/NeuralArchitectures/unet.py b/core/NeuralArchitectures/unet.py index fa1bad8..50739a4 100644 --- a/core/NeuralArchitectures/unet.py +++ b/core/NeuralArchitectures/unet.py @@ -1,58 +1,100 @@ import torch import torch.nn as nn import torch.nn.functional as F -import numpy as np -import sys -from core.NeuralLayers import * +from ..NeuralLayers import Convolution +from ..NeuralLayers import SEResidualComplex as ResSE -class BaseConv(nn.Module): - ''' - https://arxiv.org/pdf/1505.04597.pdf --> base UNet {(convolution+ReLU+batch_norm)*2} - chop: squeeze the number of channels by the ratio in the first convolution opertion + +class ConvBlock(nn.Module): + r'''Building block based on original UNet paper:: + https://arxiv.org/pdf/1505.04597.pdf + Args: + tensor_size: shape of tensor in BCHW + out_channels: depth of output feature channels + pad: True/False + norm: None/batch/group/instance/layer/pixelwise + chop: squeeze the number of channels by the ratio in the first + convolution operation ''' - def __init__(self, tensor_size, out_channels, pad=False, chop=2, *args, **kwargs): - super(BaseConv, self).__init__() + def __init__(self, tensor_size, out_channels, pad, chop=2, + *args, **kwargs): + super(ConvBlock, self).__init__() norm = "batch" + pad = pad self.baseconv = nn.Sequential() - self.baseconv.add_module("cnv_1", Convolution(tensor_size, 3, out_channels//chop, pad=pad, normalization=norm)) - self.baseconv.add_module("cnv_2", Convolution(self.baseconv[-1].tensor_size, 3, out_channels, pad=pad, normalization=norm)) + self.baseconv.add_module("cnv_1", Convolution(tensor_size, 3, + out_channels//chop, + pad=pad, + normalization=norm)) + self.baseconv.add_module("cnv_2", + Convolution(self.baseconv[-1].tensor_size, 3, + out_channels, + pad=pad, + normalization=norm)) self.tensor_size = self.baseconv[-1].tensor_size def forward(self, tensor): return self.baseconv(tensor) -class BaseResSE(nn.Module): - ''' - https://arxiv.org/pdf/1808.05238.pdf --> AnatomyNet Base Block {(residual+squeeze_excitation+leakyReLU)*2} +class ResSEBlock(nn.Module): + r''' ResSE - Residual Block with Squeeze and Excitation block: + Building block based on AnatomyNet paper:: + https://arxiv.org/pdf/1808.05238.pdf + Args: + tensor_size: shape of tensor in BCHW + out_channels: depth of output feature channels + activation: None/relu/relu6/lklu/elu/prelu/tanh/sigm/maxo/rmxo/swish + pad: True/False + norm: None/batch/group/instance/layer/pixelwise ''' def __init__(self, tensor_size, out_channels, pad=True, *args, **kwargs): - super(BaseResSE, self).__init__() - norm, activation = "batch", "lklu" - self.baseresSE = nn.Sequential() - self.baseresSE.add_module("resSE_1", SEResidualComplex(tensor_size, 3, out_channels, pad=pad, activation=activation, normalization=norm, r=4)) - self.baseresSE.add_module("resSE_2", SEResidualComplex(self.baseresSE[-1].tensor_size, 3, out_channels, pad=pad, activation=activation, normalization=norm, r=4)) - self.tensor_size = self.baseresSE[-1].tensor_size + super(ResSEBlock, self).__init__() + norm, activation, pad = "batch", "lklu", True + self.resSE = nn.Sequential() + self.resSE.add_module("resSE_1", ResSE(tensor_size, 3, out_channels, + pad=pad, activation=activation, + normalization=norm, r=4)) + self.resSE.add_module("resSE_2", ResSE(self.resSE[-1].tensor_size, 3, + out_channels, pad=pad, + activation=activation, + normalization=norm, r=4)) + self.tensor_size = self.resSE[-1].tensor_size def forward(self, tensor): - return self.baseresSE(tensor) + return self.resSE(tensor) class Down(nn.Module): - ''' + r''' Downsampling block for U-Net and AnatomyNet change: uses convolution with strides instead of pooling to downsample + Args: + tensor_size: shape of tensor in BCHW + in_channels: C in tensor_size + out_channels: depth of output feature channels + activation: None/relu/relu6/lklu/elu/prelu/tanh/sigm/maxo/rmxo/swish + pad: True/False + norm: None/batch/group/instance/layer/pixelwise ''' - def __init__(self, tensor_size, in_channels, out_channels, strides=(1,1), pad=False, dropout=0.0, nettype='anatomynet', *args, **kwargs): + def __init__(self, tensor_size, in_channels, out_channels, strides=(1, 1), + pad=False, dropout=0.0, nettype='unet', *args, **kwargs): super(Down, self).__init__() - assert nettype.lower() in ["unet", "anatomynet", "none"], "network sould be unet or anatomynet or none" + + assert nettype.lower() in ["unet", "anatomynet", "none"], \ + "network sould be unet or anatomynet or none" + self.down = nn.Sequential() - self.down.add_module("down_conv_1", Convolution(tensor_size, 2, in_channels, (2,2), pad=False)) + self.down.add_module("down1", + Convolution(tensor_size, 3, in_channels, + (2, 2), pad=True)) if nettype is 'unet': - self.down.add_module("down_conv_2", BaseConv(self.down[-1].tensor_size, out_channels, pad=pad)) + self.down.add_module("down2", + ConvBlock(self.down[-1].tensor_size, + out_channels, pad=pad)) elif nettype is 'anatomynet': - self.down.add_module("down_conv_2", BaseResSE(self.down[-1].tensor_size, out_channels, pad=pad)) - else: - pass + self.down.add_module("down2", + ResSEBlock(self.down[-1].tensor_size, + out_channels, pad=True)) self.tensor_size = self.down[-1].tensor_size def forward(self, tensor): @@ -60,105 +102,161 @@ def forward(self, tensor): class Up(nn.Module): + r'''upsampling block for U-Net and AnatomyNet + Args: + tensor_size: shape of tensor in BCHW + in_channels: C in tensor_size + out_channels: depth of output feature channels + activation: None/relu/relu6/lklu/elu/prelu/tanh/sigm/maxo/rmxo/swish + pad: True/False + norm: None/batch/group/instance/layer/pixelwise ''' - upsampling the feature maps - ''' - def __init__(self, tensor_size, out_shape, strides=(1,1), pad=False, dropout=0.0, nettype='anatomynet', *args, **kwargs): + def __init__(self, tensor_size, out_shape, strides=(2, 2), pad=False, + dropout=0.0, nettype='unet', *args, **kwargs): super(Up, self).__init__() - assert nettype.lower() in ["unet", "anatomynet", "none"], "network sould be unet or anatomynet or none" - self.up = ConvolutionTranspose(tensor_size, 2, tensor_size[1]//2, 2, False) - _ts = self.up.tensor_size - _tensor_size = (_ts[0], _ts[1]*2)+out_shape[2:] + + assert nettype.lower() in ["unet", "anatomynet", "none"],\ + "network sould be unet or anatomynet or none" + self.up = Convolution(tensor_size, 3, tensor_size[1], strides=2, + pad=True, transpose=True, maintain_out_size=True) + _tensor_size = self.up.tensor_size + self.up.tensor_size = (_tensor_size[0], _tensor_size[1]+out_shape[1]) \ + + (tensor_size[2], tensor_size[2]) + if nettype is 'unet': - self.up_base = BaseConv(_tensor_size, out_shape[1], pad=pad) + self.up_base = ConvBlock(self.up.tensor_size, out_shape[1], + pad=pad) self.tensor_size = self.up_base.tensor_size elif nettype is 'anatomynet': - self.up_base = BaseResSE(_tensor_size, out_shape[1], pad=pad) + self.up_base = ResSEBlock(self.up.tensor_size, out_shape[1], + pad=pad) self.tensor_size = self.up_base.tensor_size else: - self.tensor_size = _tensor_size - pass + self.tensor_size = self.up.tensor_size - def forward(self, tensor1,tensor2, nettype="anatomynet"): - tensor1 = self.up(tensor1, tensor2.shape) + def forward(self, tensor1, tensor2, nettype="unet"): + tensor1 = self.up(tensor1) + _, _, h, w = list(map(int.__sub__, + list(tensor1.shape), list(tensor2.shape))) + # (padLeft, padRight, padTop, padBottom) + pad = (w//2, w-w//2, h//2, h-h//2) + tensor2 = F.pad(tensor2, pad) if nettype is not "none": tensor = torch.cat([tensor2, tensor1], dim=1) return self.up_base(tensor) else: return tensor1 + class UNet(nn.Module): ''' UNet: https://arxiv.org/pdf/1505.04597.pdf 4 down blocks and 4 up blocks + Args: + tensor_size: shape of tensor in BCHW + in_channels: C in tensor_size + out_channels: initial depth of filters. + n_classes: number of output channels expected + activation: None/relu/relu6/lklu/elu/prelu/tanh/sigm/maxo/rmxo/swish + norm: None/batch/group/instance/layer/pixelwise ''' - def __init__(self, tensor_size, out_channels, n_classes = 10, *args, **kwargs): + def __init__(self, tensor_size, out_channels, n_classes, *args, **kwargs): super(UNet, self).__init__() - self.downnet1 = BaseConv(tensor_size, out_channels, pad=True) - self.downnet2 = Down(self.downnet1.tensor_size, out_channels*1, out_channels*2, pad=True) - self.downnet3 = Down(self.downnet2.tensor_size, out_channels*2, out_channels*4, pad=True) - self.downnet4 = Down(self.downnet3.tensor_size, out_channels*4, out_channels*8, pad=True) - self.downnet5 = Down(self.downnet4.tensor_size, out_channels*8, out_channels*16, pad=True) - self.upnet1 = Up(self.downnet5.tensor_size, self.downnet4.tensor_size) - self.upnet2 = Up(self.upnet1.tensor_size, self.downnet3.tensor_size) - self.upnet3 = Up(self.upnet2.tensor_size, self.downnet2.tensor_size) - self.upnet4 = Up(self.upnet3.tensor_size, self.downnet1.tensor_size) - self.final_layer= Convolution(self.upnet4.tensor_size, 1, n_classes) + out_c = out_channels + PAD = False + self.d1 = ConvBlock(tensor_size, out_c, pad=PAD) + self.d2 = Down(self.d1.tensor_size, out_c*1, out_channels*2, pad=PAD) + self.d3 = Down(self.d2.tensor_size, out_c*2, out_channels*4, pad=PAD) + self.d4 = Down(self.d3.tensor_size, out_c*4, out_channels*8, pad=PAD) + self.d5 = Down(self.d4.tensor_size, out_c*8, out_channels*16, pad=PAD) + self.u1 = Up(self.d5.tensor_size, self.d4.tensor_size) + self.u2 = Up(self.u1.tensor_size, self.d3.tensor_size) + self.u3 = Up(self.u2.tensor_size, self.d2.tensor_size) + self.u4 = Up(self.u3.tensor_size, self.d1.tensor_size) + self.final_layer = Convolution(self.u4.tensor_size, 1, n_classes) self.tensor_size = self.final_layer.tensor_size def forward(self, tensor): - d1 = self.downnet1(tensor) - d2 = self.downnet2(d1) - d3 = self.downnet3(d2) - d4 = self.downnet4(d3) - d5 = self.downnet5(d4) - u1 = self.upnet1(d5, d4) - u2 = self.upnet2(u1, d3) - u3 = self.upnet3(u2, d2) - u4 = self.upnet4(u3, d1) + d1 = self.d1(tensor) + d2 = self.d2(d1) + d3 = self.d3(d2) + d4 = self.d4(d3) + d5 = self.d5(d4) + u1 = self.u1(d5, d4) + u2 = self.u2(u1, d3) + u3 = self.u3(u2, d2) + u4 = self.u4(u3, d1) return self.final_layer(u4) -class UNetMini(nn.Module): + +class ANet(nn.Module): ''' AnatomyNet architecture -- https://arxiv.org/pdf/1808.05238.pdf + Args: + tensor_size: shape of tensor in BCHW + in_channels: C in tensor_size + out_channels: initial depth of filters. + n_classes: number of output channels expected + activation: None/relu/relu6/lklu/elu/prelu/tanh/sigm/maxo/rmxo/swish + norm: None/batch/group/instance/layer/pixelwise ''' - def __init__(self, tensor_size, out_channels, n_classes = 2, *args, **kwargs): - super(UNetMini, self).__init__() - self.downnet1 = Down(tensor_size, tensor_size[1], out_channels, pad=True) - self.basenet1 = BaseResSE(self.downnet1.tensor_size, int(out_channels*1.25)) - self.basenet2 = BaseResSE(self.basenet1.tensor_size, int(out_channels*1.50)) - self.basenet3 = BaseResSE(self.basenet2.tensor_size, int(out_channels*1.75)) - self.basenet4 = BaseResSE(self.basenet3.tensor_size, int(out_channels*1.75)) - _tensor_size = self.basenet4.tensor_size - _tensor_size = (_tensor_size[0], self.basenet4.tensor_size[1]+self.basenet2.tensor_size[1], _tensor_size[2], _tensor_size[3]) - self.concat1 = BaseResSE(_tensor_size, int(out_channels*1.50)) - _tensor_size = self.concat1.tensor_size - _tensor_size = (_tensor_size[0], self.concat1.tensor_size[1]+self.basenet1.tensor_size[1], _tensor_size[2], _tensor_size[3]) - self.concat2 = BaseResSE(_tensor_size, int(out_channels*1.25)) - _tensor_size = self.concat2.tensor_size - _tensor_size = (_tensor_size[0], self.concat2.tensor_size[1]+self.downnet1.tensor_size[1], _tensor_size[2], _tensor_size[3]) - self.concat3 = BaseResSE(_tensor_size, int(out_channels)) - self.upnet = Up(self.concat3.tensor_size, tensor_size, nettype="none") - _tensor_size = (_tensor_size[0], self.upnet.tensor_size[1]//2+tensor_size[1], _tensor_size[2], _tensor_size[3]) - self.final1 = Convolution(_tensor_size, 3, int(out_channels*0.5), pad=True, normalization="batch") - self.final2 = Convolution(self.final1.tensor_size, 3, n_classes, pad=True, normalization="batch") + def __init__(self, tensor_size, out_channels, n_classes=2, + *args, **kwargs): + super(ANet, self).__init__() + self.d1 = Down(tensor_size, tensor_size[1], out_channels, + pad=True, nettype='anatomynet') + self.b1 = ResSEBlock(self.d1.tensor_size, int(out_channels*1.25)) + self.b2 = ResSEBlock(self.b1.tensor_size, int(out_channels*1.50)) + self.b3 = ResSEBlock(self.b2.tensor_size, int(out_channels*1.75)) + self.b4 = ResSEBlock(self.b3.tensor_size, int(out_channels*1.75)) + + _tensor_size = self.b4.tensor_size + _tensor_size = (_tensor_size[0], self.b4.tensor_size[1] + + self.b2.tensor_size[1], + _tensor_size[2], _tensor_size[3]) + self.c1 = ResSEBlock(_tensor_size, int(out_channels*1.50)) + + _tensor_size = self.c1.tensor_size + _tensor_size = (_tensor_size[0], self.c1.tensor_size[1] + + self.b1.tensor_size[1], + _tensor_size[2], _tensor_size[3]) + self.c2 = ResSEBlock(_tensor_size, int(out_channels*1.25)) + + _tensor_size = self.c2.tensor_size + _tensor_size = (_tensor_size[0], self.c2.tensor_size[1] + + self.d1.tensor_size[1], + _tensor_size[2], _tensor_size[3]) + self.c3 = ResSEBlock(_tensor_size, int(out_channels)) + + self.u1 = Up(self.c3.tensor_size, tensor_size, nettype="none") + + _tensor_size = self.u1.tensor_size + self.f1 = Convolution(_tensor_size, 3, int(out_channels*0.5), pad=True, + normalization="batch") + self.f2 = Convolution(self.f1.tensor_size, 3, n_classes, pad=True, + normalization="batch") + self.tensor_size = (tensor_size[0], n_classes, + tensor_size[2], tensor_size[3]) def forward(self, tensor): - d1 = self.downnet1(tensor) - b1 = self.basenet1(d1) - b2 = self.basenet2(b1) - b3 = self.basenet3(b2) - b4 = self.basenet4(b3) - c1 = self.concat1(torch.cat([b4, b2],dim=1)) - c2 = self.concat2(torch.cat([c1, b1],dim=1)) - c3 = self.concat3(torch.cat([c2, d1],dim=1)) - up = self.upnet(c3, tensor, "none") - f1 = self.final1(torch.cat([up,tensor],dim=1)) - f2 = self.final2(f1) + assert((tensor.shape[2] % 2 == 0) & (tensor.shape[3] % 2 == 0)), \ + "tensor height and width should be divisible by 2" + d1 = self.d1(tensor) + b1 = self.b1(d1) + b2 = self.b2(b1) + b3 = self.b3(b2) + b4 = self.b4(b3) + c1 = self.c1(torch.cat([b4, b2], dim=1)) + c2 = self.c2(torch.cat([c1, b1], dim=1)) + c3 = self.c3(torch.cat([c2, d1], dim=1)) + up = self.u1(c3, tensor, "none") + f1 = self.f1(torch.cat([up, tensor], dim=1)) + f2 = self.f2(f1) return f2 -# tsize = (1,1,572,572) -# unet = UNet(tsize, 64) -# unet_mini = UNetMini(tsize, 64) + +# tsize = (1, 1, 572, 572) +# unet = UNet(tsize, 64, 2) +# anet = ANet(tsize, 32) # unet(torch.rand(tsize)).shape -# unet_mini(torch.rand(tsize)).shape +# anet(torch.rand(tsize)).shape diff --git a/core/NeuralLayers/__init__.py b/core/NeuralLayers/__init__.py index d369ec5..82b1db9 100644 --- a/core/NeuralLayers/__init__.py +++ b/core/NeuralLayers/__init__.py @@ -24,8 +24,7 @@ from .sae import ConvolutionalSAE from .detailpooling import DetailPooling -from .lossfunctions import CapsuleLoss, CategoricalLoss, TripletLoss, \ - DiceLoss +from .lossfunctions import CapsuleLoss, CategoricalLoss, TripletLoss, DiceLoss from .obfuscatedecolor import ObfuscateDecolor from .activations import Activations diff --git a/core/NeuralLayers/lossfunctions.py b/core/NeuralLayers/lossfunctions.py index 668dc71..a48940a 100644 --- a/core/NeuralLayers/lossfunctions.py +++ b/core/NeuralLayers/lossfunctions.py @@ -43,31 +43,45 @@ def nlog_likelihood(tensor, targets): return F.nll_loss(tensor.log_softmax(1), targets) # =========================================================================== # -def hardest_negative(lossValues,margin): + +def hardest_negative(lossValues, margin): return lossValues.max(2)[0].max(1)[0].mean() + def semihard_negative(lossValues, margin): - lossValues = torch.where((torch.ByteTensor(lossValues>0.) & torch.ByteTensor(lossValues 0.) & + torch.ByteTensor(lossValues < margin)), + lossValues, torch.zeros(lossValues.size())) return lossValues.max(2)[0].max(1)[0].mean() class TripletLoss(nn.Module): - def __init__(self, margin, negative_selection_fn='hardest_negative', samples_per_class = 2, *args, **kwargs): + def __init__(self, margin, negative_selection_fn='hardest_negative', + samples_per_class=2, *args, **kwargs): super(TripletLoss, self).__init__() self.tensor_size = (1,) self.margin = margin self.negative_selection_fn = negative_selection_fn - self.sqrEuc = lambda x : (x.unsqueeze(0) - x.unsqueeze(1)).pow(2).sum(2).div(x.size(1)) + self.sqrEuc = lambda x: (x.unsqueeze(0) - + x.unsqueeze(1)).pow(2).sum(2).div(x.size(1)) self.perclass = samples_per_class def forward(self, embeddings, labels): - InClass = labels.reshape(-1,1) == labels.reshape(1,-1) - Consider = torch.eye(labels.size(0)).mul(-1).add(1).type(InClass.type()) - Scores = self.sqrEuc(embeddings) - Gs = Scores.view(-1, 1)[(InClass*Consider).view(-1, 1)].reshape(-1, self.perclass-1) - Is = Scores.view(-1, 1)[(InClass == 0).view(-1, 1)].reshape(-1, embeddings.size(0)-self.perclass) - lossValues = Gs.view(embeddings.size(0), -1, 1) - Is.view(embeddings.size(0), 1, -1) + self.margin + labels = torch.from_numpy(np.array([1, 1, 0, 1, 1], dtype='float32')) + InClass = labels.reshape(-1, 1) == labels.reshape(1, -1) + Consider = torch.eye(labels.size(0)).mul(-1).add(1) \ + .type(InClass.type()) + Scores = self.sqrEuc(embeddings) + + Gs = Scores.view(-1, 1)[(InClass*Consider).view(-1, 1)] \ + .reshape(-1, self.perclass-1) + Is = Scores.view(-1, 1)[(InClass == 0).view(-1, 1)] \ + .reshape(-1, embeddings.size(0)-self.perclass) + + lossValues = Gs.view(embeddings.size(0), -1, 1) - \ + Is.view(embeddings.size(0), 1, -1) + self.margin lossValues = lossValues.clamp(0.) + if self.negative_selection_fn == "hardest_negative": return hardest_negative(lossValues, self.margin), Gs, Is elif self.negative_selection_fn == "semihard_negative": @@ -78,15 +92,22 @@ def forward(self, embeddings, labels): class DiceLoss(nn.Module): - """ + r""" Dice/ Tversky loss for semantic segmentationself. Implemented from https://arxiv.org/pdf/1803.11078.pdf - https://arxiv.org/pdf/1706.05721.pdf has same equation but with alpha and beta controlling FP and FN. - works for both softmax predicted output [batchx2xH,W] or sigmoid output [batchx1xH,W - p_i * g_i --> True Positives (TP) - p_i * g_j --> False Positives (FP) - p_j * g_i --> False Negatives (FN) + https://arxiv.org/pdf/1706.05721.pdf has same equation but with alpha + and beta controlling FP and FN. + Args: + type: tversky/dice + Definations: + p_i - correctly predicted foreground pixels + p_j - correctly predicted background pixels + g_i - target foreground pixels + g_j - target background pixels + p_i * g_i - True Positives (TP) + p_i * g_j - False Positives (FP) + p_j * g_i - False Negatives (FN) """ - def __init__(self, type = "tversky", *args, **kwargs): + def __init__(self, type="tversky", *args, **kwargs): super(DiceLoss, self).__init__() self.tensor_size = (1,) if type == "tversky": @@ -97,6 +118,7 @@ def __init__(self, type = "tversky", *args, **kwargs): self.beta = 1.0 # below Eq(6) else: raise NotImplementedError + def forward(self, prediction, targets): top1, top5 = 0., 0. if prediction.shape[1] == 1: @@ -106,21 +128,23 @@ def forward(self, prediction, targets): g_j = targets.mul(-1).add(1) # the above is similar to one hot encoding of targets num = (p_i*g_i).sum(1).sum(1).mul((1 + self.beta**2)) # eq(5) - den = num.add((p_i*g_j).sum(1).sum(1).mul((self.beta**2))).add((p_j*g_i).sum(1).sum(1).mul((self.beta))) # eq(5) + den = num.add((p_i*g_j).sum(1).sum(1).mul((self.beta**2))) \ + .add((p_j*g_i).sum(1).sum(1).mul((self.beta))) # eq(5) loss = num / den.add(1e-6) elif prediction.shape[1] == 2: - p_i = prediction[:,0,:,:] - p_j = prediction[:,1,:,:] + p_i = prediction[:, 0, :, :] + p_j = prediction[:, 1, :, :] g_i = targets g_j = targets.mul(-1).add(1) # the above is similar to one hot encoding of targets num = (p_i*g_i).sum(1).sum(1).mul((1 + self.beta**2)) # eq(5) - den = num.add((p_i*g_j).sum(1).sum(1).mul((self.beta**2))).add((p_j*g_i).sum(1).sum(1).mul((self.beta))) # eq(5) + den = num.add((p_i*g_j).sum(1).sum(1).mul((self.beta**2))) \ + .add((p_j*g_i).sum(1).sum(1).mul((self.beta))) # eq(5) loss = num / den.add(1e-6) else: raise NotImplementedError return loss.mean(), (top1, top5) -# ============================================================================ # +# =========================================================================== # class CapsuleLoss(nn.Module):