def __init__(self,
inchannel,
channel,
stride=1,
downsample=None,
normal='bn'):
super(BasicBlock, self).__init__()
self.downsample = downsample
self.conv1 = conv3x3(inchannel,
channel,
padding=1,
stride=stride,
normal=normal)
self.conv2 = conv3x3(channel,
channel * self.expansion,
padding=1,
normal=None,
rl=False)
if normal == 'bn':
self.normal = nn.BatchNorm2d(channel * self.expansion)
elif normal == 'gn':
self.normal = nn.GroupNorm(group_num, channel * self.expansion)
self.relu = nn.ReLU()
def _make_layer(self,
block,
channel_num,
block_num,
change_size=False,
normal='bn'):
# change_size=True means the first block need change input size
if change_size:
stride = 2
else:
stride = 1
if block is BasicBlock:
if change_size:
downsample = conv3x3(self.inchannel_num,
channel_num * block.expansion,
padding=1,
stride=stride,
normal=normal)
else:
downsample = None
elif block is Bottleneck:
downsample = conv1x1(self.inchannel_num,
channel_num * block.expansion,
stride=stride,
normal=normal)
else:
raise ValueError('"block" should be "BasicBlock" or "Bottleneck"')
layers = list()
layers.append(
block(self.inchannel_num,
channel_num,
stride=stride,
downsample=downsample,
normal=normal))
for _ in range(1, block_num):
layers.append(
block(channel_num * block.expansion,
channel_num,
normal=normal))
self.inchannel_num = channel_num * block.expansion
return nn.Sequential(*layers)
def __init__(self, in_channel=3, class_num=2, normal='bn'):
super(UNetv1, self).__init__()
# down sample stage
self.conv_1 = nn.Sequential(conv3x3(in_channel, 64, padding=1, normal=normal),
conv3x3(64, 64, padding=1, normal=normal))
self.conv_2 = nn.Sequential(conv3x3(64, 128, padding=1, normal=normal),
conv3x3(128, 128, padding=1, normal=normal))
self.conv_3 = nn.Sequential(conv3x3(128, 256, padding=1, normal=normal),
conv3x3(256, 256, padding=1, normal=normal))
self.conv_4 = nn.Sequential(conv3x3(256, 512, padding=1, normal=normal),
conv3x3(512, 512, padding=1, normal=normal))
self.conv_5 = nn.Sequential(conv3x3(512, 1024, padding=1, normal=normal),
conv3x3(1024, 1024, padding=1, normal=normal))
self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
# upsample stage
# up_conv_4 corresponds conv_4
self.up_conv_4 = nn.Sequential(up_conv2x2(1024, 512))
# conv the cat(stage_4,up_conv_4) from 1024 to 512
self.conv_6 = nn.Sequential(conv3x3(1024, 512, padding=1, normal=normal),
conv3x3(512, 512, padding=1, normal=normal))
# up_conv_3 corresponds conv_3
self.up_conv_3 = nn.Sequential(up_conv2x2(512, 256))
# conv the cat(stage_3,up_conv_3) from 512 to 256
self.conv_7 = nn.Sequential(conv3x3(512, 256, padding=1, normal=normal),
conv3x3(256, 256, padding=1, normal=normal))
# up_conv_2 corresponds conv_2
self.up_conv_2 = nn.Sequential(up_conv2x2(256,128))
# conv the cat(stage_2,up_conv_2) from 256 to 128
self.conv_8 = nn.Sequential(conv3x3(256, 128, padding=1, normal=normal),
conv3x3(128, 128, padding=1, normal=normal))
# up_conv_1 corresponds conv_1
self.up_conv_1 = nn.Sequential(up_conv2x2(128, 64))
# conv the cat(stage_1,up_conv_1) from 128 to 64
self.conv_9 = nn.Sequential(conv3x3(128, 64, padding=1, normal=normal),
conv3x3(64, 64, padding=1, normal=normal))
# output
self.result = conv1x1(64, class_num, normal=normal)
def __init__(self, in_channel=3, class_num=2):
super(UNet, self).__init__()
self.conv_1 = nn.Sequential(conv3x3(in_channel, 64), conv3x3(64, 64))
self.conv_2 = nn.Sequential(conv3x3(64, 128), conv3x3(128, 128))
self.conv_3 = nn.Sequential(conv3x3(128, 256), conv3x3(256, 256))
self.conv_4 = nn.Sequential(conv3x3(256, 512), conv3x3(512, 512))
self.conv_5 = nn.Sequential(conv3x3(512, 1024), conv3x3(1024, 1024))
self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
self.up_conv_4 = nn.Sequential(up_conv2x2(1024, 512))
self.conv_6 = nn.Sequential(conv3x3(1024, 512), conv3x3(512, 512))
self.up_conv_3 = nn.Sequential(up_conv2x2(512, 256))
self.conv_7 = nn.Sequential(conv3x3(512, 256), conv3x3(256, 256))
self.up_conv_2 = nn.Sequential(up_conv2x2(256, 128))
self.conv_8 = nn.Sequential(conv3x3(256, 128), conv3x3(128, 128))
self.up_conv_1 = nn.Sequential(up_conv2x2(128, 64))
self.conv_9 = nn.Sequential(conv3x3(128, 64), conv3x3(64, 64))
self.result = conv1x1(64, class_num)