def __init__(self, params):
super(UNet2D, self).__init__()
self.params = params
self.in_chns = self.params['in_chns']
self.ft_chns = self.params['feature_chns']
self.n_class = self.params['class_num']
self.acti_func = self.params['acti_func']
assert(len(self.ft_chns) == 5)
self.block1 = UNetBlock(self.in_chns, self.ft_chns[0],
self.acti_func, self.params)
self.block2 = UNetBlock(self.ft_chns[0], self.ft_chns[1],
self.acti_func, self.params)
self.block3 = UNetBlock(self.ft_chns[1], self.ft_chns[2],
self.acti_func, self.params)
self.block4 = UNetBlock(self.ft_chns[2], self.ft_chns[3],
self.acti_func, self.params)
self.block5 = UNetBlock(self.ft_chns[3], self.ft_chns[4],
self.acti_func, self.params)
self.block6 = UNetBlock(self.ft_chns[3] * 2, self.ft_chns[3],
self.acti_func, self.params)
self.block7 = UNetBlock(self.ft_chns[2] * 2, self.ft_chns[2],
self.acti_func, self.params)
self.block8 = UNetBlock(self.ft_chns[1] * 2, self.ft_chns[1],
self.acti_func, self.params)
self.block9 = UNetBlock(self.ft_chns[0] * 2, self.ft_chns[0],
self.acti_func, self.params)
self.down1 = nn.MaxPool2d(kernel_size = 2)
self.down2 = nn.MaxPool2d(kernel_size = 2)
self.down3 = nn.MaxPool2d(kernel_size = 2)
self.down4 = nn.MaxPool2d(kernel_size = 2)
self.up1 = DeconvolutionLayer(self.ft_chns[4], self.ft_chns[3], kernel_size = 2,
dim = 2, stride = 2, acti_func = get_acti_func(self.acti_func, self.params))
self.up2 = DeconvolutionLayer(self.ft_chns[3], self.ft_chns[2], kernel_size = 2,
dim = 2, stride = 2, acti_func = get_acti_func(self.acti_func, self.params))
self.up3 = DeconvolutionLayer(self.ft_chns[2], self.ft_chns[1], kernel_size = 2,
dim = 2, stride = 2, acti_func = get_acti_func(self.acti_func, self.params))
self.up4 = DeconvolutionLayer(self.ft_chns[1], self.ft_chns[0], kernel_size = 2,
dim = 2, stride = 2, acti_func = get_acti_func(self.acti_func, self.params))
self.conv = nn.Conv2d(self.ft_chns[0], self.n_class,
kernel_size = 3, padding = 1)
def __init__(self, params):
super(MyUNet2D, self).__init__()
self.params = params
self.in_chns = self.params['in_chns']
self.ft_chns = self.params['feature_chns']
self.n_class = self.params['class_num']
self.acti_func = self.params['acti_func']
self.dropout = self.params['dropout']
self.resolution_level = len(self.ft_chns)
assert(self.resolution_level == 4)
self.block1 = MyUNetBlock(self.in_chns, self.ft_chns[0],
self.acti_func, self.params)
self.block2 = MyUNetBlock(self.ft_chns[0], self.ft_chns[1],
self.acti_func, self.params)
self.block3 = MyUNetBlock(self.ft_chns[1], self.ft_chns[2],
self.acti_func, self.params)
self.block4 = MyUNetBlock(self.ft_chns[2], self.ft_chns[3],
self.acti_func, self.params)
self.block5 = MyUNetBlock(self.ft_chns[2] * 2, self.ft_chns[2],
self.acti_func, self.params)
self.block6 = MyUNetBlock(self.ft_chns[1] * 2, self.ft_chns[1],
self.acti_func, self.params)
self.block7 = MyUNetBlock(self.ft_chns[0] * 2, self.ft_chns[0],
self.acti_func, self.params)
self.down1 = nn.MaxPool2d(kernel_size = 2)
self.down2 = nn.MaxPool2d(kernel_size = 2)
self.down3 = nn.MaxPool2d(kernel_size = 2)
self.up1 = DeconvolutionLayer(self.ft_chns[3], self.ft_chns[2], kernel_size = 2,
dim = 2, stride = 2, acti_func = get_acti_func(self.acti_func, self.params))
self.up2 = DeconvolutionLayer(self.ft_chns[2], self.ft_chns[1], kernel_size = 2,
dim = 2, stride = 2, acti_func = get_acti_func(self.acti_func, self.params))
self.up3 = DeconvolutionLayer(self.ft_chns[1], self.ft_chns[0], kernel_size = 2,
dim = 2, stride = 2, acti_func = get_acti_func(self.acti_func, self.params))
if(self.dropout):
self.drop3 = nn.Dropout(0.3)
self.drop4 = nn.Dropout(0.5)
self.conv = nn.Conv2d(self.ft_chns[0], self.n_class,
kernel_size = 3, padding = 1)
def __init__(self, in_size, out_size, is_batchnorm=True):
super(UnetUp3_CT2, self).__init__()
self.conv = UNetBlock(in_size + out_size,
out_size,
is_batchnorm,
kernel_size=(1, 3, 3),
padding_size=(0, 1, 1))
self.up = DeconvolutionLayer(kernel_size=(1, 2, 2),
stride=(1, 2, 2),
acti_func=get_acti_func(
self.acti_func, self.params))
def __init__(self, in_channels, out_chnannels,
dim, resample, acti_func, acti_func_param):
super(UNetBlock, self).__init__()
self.in_chns = in_channels
self.out_chns = out_chnannels
self.dim = dim
self.resample = resample # resample should be 'down', 'up', or None
self.acti_func = acti_func
self.conv1 = ConvolutionLayer(in_channels, out_chnannels, kernel_size = 3, padding=1,
dim = self.dim, acti_func=get_acti_func(acti_func, acti_func_param))
self.conv2 = ConvolutionLayer(out_chnannels, out_chnannels, kernel_size = 3, padding=1,
dim = self.dim, acti_func=get_acti_func(acti_func, acti_func_param))
if(self.resample == 'down'):
if(self.dim == 2):
self.resample_layer = nn.MaxPool2d(kernel_size = 2, stride = 2)
else:
self.resample_layer = nn.MaxPool3d(kernel_size = 2, stride = 2)
elif(self.resample == 'up'):
self.resample_layer = DeconvolutionLayer(out_chnannels, out_chnannels, kernel_size = 2,
dim = self.dim, stride = 2, acti_func = get_acti_func(acti_func, acti_func_param))
else:
assert(self.resample == None)
def __init__(self, params):
super(Baseunet2d5_att_pe, self).__init__()
self.params = params
self.in_chns = self.params['in_chns']
self.ft_chns = self.params['feature_chns']
self.n_class = self.params['class_num']
self.acti_func = self.params['acti_func']
self.dropout = self.params['dropout']
assert (len(self.ft_chns) == 5)
self.block1 = UNetBlock(self.in_chns, self.ft_chns[0], (1, 3, 3),
(0, 1, 1), self.acti_func, self.params)
self.block2 = UNetBlock(self.ft_chns[0], self.ft_chns[1], (1, 3, 3),
(0, 1, 1), self.acti_func, self.params)
self.block3 = UNetBlock(self.ft_chns[1], self.ft_chns[2], (1, 3, 3),
(0, 1, 1), self.acti_func, self.params)
self.block4 = UNetBlock(self.ft_chns[2], self.ft_chns[3], (1, 3, 3),
(0, 1, 1), self.acti_func, self.params)
self.block5 = UNetBlock(self.ft_chns[3], self.ft_chns[4], (3, 3, 3),
(1, 1, 1), self.acti_func, self.params)
self.block6 = UNetBlock(self.ft_chns[3] * 2, self.ft_chns[3],
(1, 3, 3), (0, 1, 1), self.acti_func,
self.params)
self.block7 = UNetBlock(self.ft_chns[2] * 2, self.ft_chns[2],
(1, 3, 3), (0, 1, 1), self.acti_func,
self.params)
self.block8 = UNetBlock(self.ft_chns[1] * 2, self.ft_chns[1],
(1, 3, 3), (0, 1, 1), self.acti_func,
self.params)
self.block9 = UNetBlock(self.ft_chns[0] * 2, self.ft_chns[0],
(1, 3, 3), (0, 1, 1), self.acti_func,
self.params)
self.down1 = nn.MaxPool3d(kernel_size=(1, 2, 2), stride=(1, 2, 2))
self.down2 = nn.MaxPool3d(kernel_size=(1, 2, 2), stride=(1, 2, 2))
self.down3 = nn.MaxPool3d(kernel_size=(1, 2, 2), stride=(1, 2, 2))
self.down4 = nn.MaxPool3d(kernel_size=2)
self.up1 = DeconvolutionLayer(self.ft_chns[4],
self.ft_chns[3],
kernel_size=1,
stride=1,
acti_func=get_acti_func(
self.acti_func, self.params))
self.up2 = DeconvolutionLayer(self.ft_chns[3],
self.ft_chns[2],
kernel_size=(1, 2, 2),
stride=(1, 2, 2),
acti_func=get_acti_func(
self.acti_func, self.params))
self.up3 = DeconvolutionLayer(self.ft_chns[2],
self.ft_chns[1],
kernel_size=(1, 2, 2),
stride=(1, 2, 2),
acti_func=get_acti_func(
self.acti_func, self.params))
self.up4 = DeconvolutionLayer(self.ft_chns[1],
self.ft_chns[0],
kernel_size=(1, 2, 2),
stride=(1, 2, 2),
acti_func=get_acti_func(
self.acti_func, self.params))
self.conv = nn.Conv3d(self.ft_chns[0],
self.n_class,
kernel_size=(1, 3, 3),
padding=(0, 1, 1))
self.pe1 = ProjectExciteLayer(self.ft_chns[0])
self.pe2 = ProjectExciteLayer(self.ft_chns[1])
self.pe3 = ProjectExciteLayer(self.ft_chns[2])
self.pe4 = ProjectExciteLayer(self.ft_chns[3])
self.pe5 = ProjectExciteLayer(self.ft_chns[3] * 2)
self.pe6 = ProjectExciteLayer(self.ft_chns[2] * 2)
self.pe7 = ProjectExciteLayer(self.ft_chns[1] * 2)
self.pe8 = ProjectExciteLayer(self.ft_chns[0] * 2)
# dropout
if (self.dropout):
self.drop1 = nn.Dropout(p=0)
self.drop2 = nn.Dropout(p=0)
self.drop3 = nn.Dropout(p=0)
self.drop4 = nn.Dropout(p=0.3)
self.drop5 = nn.Dropout(p=0.5)
#attention block(two gates)
self.attention1 = AttentionBlock(self.ft_chns[3] * 2, self.ft_chns[3])
self.attention2 = AttentionBlock(self.ft_chns[2] * 2, self.ft_chns[2])
self.attention3 = AttentionBlock(self.ft_chns[1] * 2, self.ft_chns[1])
self.attention4 = AttentionBlock(self.ft_chns[0] * 2, self.ft_chns[0])
self.attention5 = AttentionBlock(self.ft_chns[3], self.ft_chns[3] // 2)
def __init__(self, params):
super(UNet3D, self).__init__()
self.params = params
self.in_chns = self.params['in_chns']
self.ft_chns = self.params['feature_chns']
self.n_class = self.params['class_num']
self.acti_func = self.params['acti_func']
self.dropout = self.params['dropout']
self.resolution_level = len(self.ft_chns)
assert (self.resolution_level == 5 or self.resolution_level == 4)
self.block1 = UNetBlock(self.in_chns, self.ft_chns[0], self.acti_func,
self.params)
self.block2 = UNetBlock(self.ft_chns[0], self.ft_chns[1],
self.acti_func, self.params)
self.block3 = UNetBlock(self.ft_chns[1], self.ft_chns[2],
self.acti_func, self.params)
self.block4 = UNetBlock(self.ft_chns[2], self.ft_chns[3],
self.acti_func, self.params)
if (self.resolution_level == 5):
self.block5 = UNetBlock(self.ft_chns[3], self.ft_chns[4],
self.acti_func, self.params)
self.block6 = UNetBlock(self.ft_chns[3] * 2, self.ft_chns[3],
self.acti_func, self.params)
self.block7 = UNetBlock(self.ft_chns[2] * 2, self.ft_chns[2],
self.acti_func, self.params)
self.block8 = UNetBlock(self.ft_chns[1] * 2, self.ft_chns[1],
self.acti_func, self.params)
self.block9 = UNetBlock(self.ft_chns[0] * 2, self.ft_chns[0],
self.acti_func, self.params)
self.down1 = nn.MaxPool3d(kernel_size=2)
self.down2 = nn.MaxPool3d(kernel_size=2)
self.down3 = nn.MaxPool3d(kernel_size=2)
if (self.resolution_level == 5):
self.down4 = nn.MaxPool3d(kernel_size=2)
self.up1 = DeconvolutionLayer(self.ft_chns[4],
self.ft_chns[3],
kernel_size=2,
stride=2,
acti_func=get_acti_func(
self.acti_func, self.params))
self.up2 = DeconvolutionLayer(self.ft_chns[3],
self.ft_chns[2],
kernel_size=2,
stride=2,
acti_func=get_acti_func(
self.acti_func, self.params))
self.up3 = DeconvolutionLayer(self.ft_chns[2],
self.ft_chns[1],
kernel_size=2,
stride=2,
acti_func=get_acti_func(
self.acti_func, self.params))
self.up4 = DeconvolutionLayer(self.ft_chns[1],
self.ft_chns[0],
kernel_size=2,
stride=2,
acti_func=get_acti_func(
self.acti_func, self.params))
if (self.dropout):
self.drop1 = nn.Dropout(p=0.1)
self.drop2 = nn.Dropout(p=0.1)
self.drop3 = nn.Dropout(p=0.2)
self.drop4 = nn.Dropout(p=0.2)
if (self.resolution_level == 5):
self.drop5 = nn.Dropout(p=0.3)
self.conv = nn.Conv3d(self.ft_chns[0],
self.n_class,
kernel_size=3,
padding=1)