def __init__(self, in_channels, BN, bn_eps=1e-5):
"""
Due to this inception block contains sub inception blocks, we have to define it independently.
"""
super(InceptionE, self).__init__()
self.branch1x1 = omth_blocks.conv_block(in_channels,
filters=[320],
kernel_sizes=[1],
stride=[1],
padding=[0],
batch_norm=BN)
self.branch3x3_1 = omth_blocks.conv_block(in_channels,
filters=[384],
kernel_sizes=[1],
stride=[1],
padding=[0],
batch_norm=BN)
self.branch3x3_2 = self.sub_inception_module(BN)
self.branch3x3dbl_1 = omth_blocks.conv_block(in_channels,
filters=[384, 384],
kernel_sizes=[1, 3],
stride=[1, 1],
padding=[0, 1],
batch_norm=BN)
self.branch3x3dbl_2 = self.sub_inception_module(BN)
self.branch_pool = omth_blocks.conv_block(in_channels,
filters=[192],
kernel_sizes=[1],
stride=[1],
padding=[0],
batch_norm=BN)
def create_backbone_model(self):
# Prepare VGG-16 net with batch normalization
vgg16_model = vgg16_bn(pretrained=True)
net = list(vgg16_model.children())[0]
# Replace the maxout with ceil in vanilla vgg16 net
ceil_maxout = nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=True)
net = [ceil_maxout if type(n) is nn.MaxPool2d else n for n in net]
# Basic VGG Layers
self.conv_module_name.append("conv_1")
self.conv_module.append(nn.Sequential(*net[:6]))
self.conv_module_name.append("conv_2")
# block = omth_blocks.conv_block(64, filters=[64, 128],kernel_sizes=[3, 3], stride=[1, 1],
# padding=[1, 1],batch_norm=batch_norm)
# block.add_module("2", ceil_maxout)
# self.conv_module.append(block)
self.conv_module.append(nn.Sequential(*net[6:13]))
self.conv_module_name.append("conv_3")
# block = omth_blocks.conv_block(128, filters=[128, 128], kernel_sizes=[3, 1], stride=[1, 1],
# padding=[1, 0], batch_norm=batch_norm)
# block.add_module("3", ceil_maxout)
# self.conv_module.append(block)
self.conv_module.append(nn.Sequential(*net[13:23]))
self.conv_module_name.append("conv_4")
# block = omth_blocks.conv_block(128, filters=[256, 256], kernel_sizes=[3, 1], stride=[1, 1],
# padding=[1, 0], batch_norm=batch_norm)
# block.add_module("4", ceil_maxout)
# self.conv_module.append(block)
self.conv_module.append(nn.Sequential(*net[23:33]))
self.conv_module_name.append("conv_5")
# block = omth_blocks.conv_block(256, filters=[384, 384], kernel_sizes=[3, 1], stride=[1, 1],
# padding=[1, 0], batch_norm=batch_norm)
# block.add_module("5", ceil_maxout)
# self.conv_module.append(block)
self.conv_module.append(nn.Sequential(*net[33:43]))
# Extra Layers
self.conv_module_name.append("extra_1")
self.conv_module.append(omth_blocks.conv_block(self.bottleneck_channel, [512, 512],
kernel_sizes=[3, 1], stride=[1, 1], padding=[3, 0],
dilation=[3, 1], batch_norm=nn.BatchNorm2d))
self.conv_module_name.append("extra_2")
self.conv_module.append(omth_blocks.conv_block(512, [256, 512], kernel_sizes=[1, 3],
stride=[1, 2], padding=[0, 1], batch_norm=nn.BatchNorm2d))
self.conv_module_name.append("extra_3")
self.conv_module.append(omth_blocks.conv_block(512, [256, 256], kernel_sizes=[1, 3],
stride=[1, 1], padding=[0, 1], batch_norm=nn.BatchNorm2d))
self.conv_module_name.append("extra_4")
self.conv_module.append(omth_blocks.conv_block(256, [256, 256], kernel_sizes=[1, 3],
stride=[1, 1], padding=[0, 1], batch_norm=nn.BatchNorm2d))
def create_loc_layer(self,
in_channel,
anchor,
stride,
incep_loc=False,
in_wid=128):
loc_layer = nn.ModuleList([])
loc_layer.append(
omth_blocks.conv_block(in_channel, [in_channel, in_channel],
kernel_sizes=[3, 1],
stride=[1, 1],
padding=[3, 0],
dilation=[3, 1],
batch_norm=self.batch_norm))
if incep_loc:
loc_layer.append(
omth_blocks.InceptionBlock(in_channel,
filters=[[128, 128, 128, in_wid],
[128, 128, 128, in_wid],
[128, 128, in_wid],
[192, 192, 128, in_wid]],
kernel_sizes=[[[1, 9], [1, 5], 3,
1],
[[1, 7], [1, 3], 3,
1], [[1, 5], 3, 1],
[[1, 3], [3, 1], 3,
1]],
stride=[[1, 1, 1, 1], [1, 1, 1, 1],
[1, 1, 1], [1, 1, 1, 1]],
padding=[[[0, 4], [0, 2], 1, 0],
[[0, 3], [0, 1], 1, 0],
[[0, 2], 1, 0],
[[0, 1], [1, 0], 1, 0]],
batch_norm=None,
inner_maxout=None))
input_channel = in_wid * 4 if incep_loc else in_channel
loc_layer.append(
omth_blocks.conv_block(
input_channel,
filters=[input_channel,
int(input_channel / 2), anchor * 4],
kernel_sizes=[3, 1, 3],
stride=[1, 1, stride],
padding=[0, 1, 1],
activation=None))
loc_layer.apply(init_cnn)
return loc_layer
def create_conf_layer(self, in_channel, anchor, stride, incep_conf=False):
conf_layer = nn.ModuleList([])
conf_layer.append(omth_blocks.conv_block(
in_channel, [in_channel, in_channel], kernel_sizes=[3, 1], stride=[1, 1],
padding=[3, 0], dilation=[3, 1], batch_norm=self.batch_norm)
)
if self.connect_loc_to_conf:
if incep_conf:
out_chnl = int(in_channel / 8)
out_chnl_2 = int(in_channel / 2) - (3 * out_chnl)
conf_layer.append(omth_blocks.InceptionBlock(
in_channel, stride=[[1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1]],
kernel_sizes=[[[1, 7], 3, 1], [[1, 5], 3, 1], [[1, 3], 3, 1], [3, 1]],
filters=[[64, 64, out_chnl], [64, 64, out_chnl], [64, 64, out_chnl], [128, out_chnl_2]],
padding=[[[0, 3], 1, 0], [[0, 2], 1, 0], [[0, 1], 1, 0], [1, 0]],
batch_norm=None, inner_maxout=None))
else:
conf_layer.append(omth_blocks.conv_block(
in_channel, filters=[in_channel, int(in_channel / 2)],
kernel_sizes=[1, 3], stride=[1, 1], padding=[0, 1], activation=None))
# In this layer, the output from loc_layer will be concatenated to the conf layer
# Feeding the conf layer with regressed location, helping the conf layer
# to get better prediction
conf_concate = omth_blocks.conv_block(
int(in_channel / 2) + anchor * 4, kernel_sizes=[3, 1, 3],
filters=[int(in_channel / 2), int(in_channel / 4), anchor * 2],
stride=[1, 1, stride], padding=[1, 0, 1], activation=None)
conf_concate.apply(init.init_cnn)
else:
print("incep_conf is turned off due to connect_loc_to_conf is False")
conf_layer.append(omth_blocks.conv_block(
in_channel, filters=[in_channel, int(in_channel / 2), anchor * 2],
kernel_sizes=[1, 3, 3], stride=[1, 1, stride], padding=[0, 1, 1], activation=None))
conf_concate = None
conf_layer.apply(init.init_cnn)
return conf_layer, conf_concate
def __init__(self, in_channel, filters, final_inchannel, activation, depth,
out_depth):
super().__init__()
self.sub_module = nn.ModuleList([])
for _ in range(depth):
if _ == depth - 1:
filter = out_depth
in_channel = final_inchannel
else:
filter = filters
self.sub_module.append(
# Naive Block
block.conv_block(in_channel,
filters=[filter, filter],
kernel_sizes=[3, 1],
stride=[1, 1],
padding=[1, 0],
activation=activation))
"""
def __init__(self,
input_channel,
output_channel,
BN=nn.BatchNorm2d,
upscale_factor=2):
super().__init__()
"""
module = []
module.append(nn.Conv2d(in_channels=input_channel, kernel_size=3, padding=1,
out_channels=input_channel * upscale_factor * upscale_factor))
module.append(nn.PixelShuffle(upscale_factor))
self.up_conv = nn.Sequential(*module)
self.norm_conv = omth_blocks.conv_block(input_channel, kernel_sizes=1,
filters=output_channel, stride=1, padding=0,
batch_norm=BN)
"""
self.up_conv = omth_blocks.conv_block(
input_channel,
kernel_sizes=[2 + upscale_factor, 1],
filters=[output_channel, output_channel],
stride=[0 + upscale_factor, 1],
padding=[1, 0],
batch_norm=BN,
transpose=[True, False])
def __init__(self, BN, num_classes=1000, aux_logits=True):
"""
:param BN:
:param num_classes:
:param aux_logits:
"""
super().__init__()
self.avg_pool_3_1_1 = nn.AvgPool2d(kernel_size=3, stride=1, padding=1)
self.pool_3_2_0 = nn.MaxPool2d(kernel_size=3, stride=2, padding=0)
self.pool_8_1_0 = nn.MaxPool2d(kernel_size=8, stride=1, padding=0)
self.BN = BN
self.aux_logits = aux_logits
self.conv_layers_1 = omth_blocks.conv_block(3,
filters=[32, 32, 64],
kernel_sizes=[3, 3, 3],
stride=[2, 1, 1],
padding=[0, 0, 1],
name="conv_block_1",
batch_norm=BN)
self.conv_layers_2 = omth_blocks.conv_block(64,
filters=[80, 192],
kernel_sizes=[1, 3],
stride=[1, 1],
padding=[0, 0],
name="conv_block_2",
batch_norm=BN)
self.Mixed5b = self.inceptionA(in_channels=192,
pool_features=32,
pool=self.avg_pool_3_1_1)
self.Mixed5c = self.inceptionA(in_channels=256,
pool_features=64,
pool=self.avg_pool_3_1_1)
self.Mixed5d = self.inceptionA(in_channels=288,
pool_features=64,
pool=self.avg_pool_3_1_1)
self.Mixed6a = omth_blocks.InceptionBlock(288,
filters=[[384], [64, 96,
96]],
kernel_sizes=[[3], [1, 3,
3]],
stride=[[2], [1, 1, 2]],
padding=[[0], [0, 1, 0]],
batch_norm=BN,
maxout=self.pool_3_2_0)
self.Mixed6b = self.inceptionC(in_channels=768,
c7=120,
pool=self.avg_pool_3_1_1)
self.Mixed6c = self.inceptionC(in_channels=768,
c7=160,
pool=self.avg_pool_3_1_1)
self.Mixed6d = self.inceptionC(in_channels=768,
c7=160,
pool=self.avg_pool_3_1_1)
self.Mixed6e = self.inceptionC(in_channels=768,
c7=192,
pool=self.avg_pool_3_1_1)
self.Mixed7a = omth_blocks.InceptionBlock(
768,
filters=[[192, 320], [192, 192, 192, 192]],
kernel_sizes=[[1, 3], [1, [7, 1], [1, 7], 3]],
stride=[[1, 2], [1, 1, 1, 2]],
padding=[[0, 0], [0, [0, 3], [3, 0], 0]],
maxout=self.pool_3_2_0,
batch_norm=BN)
self.Mixed7b = InceptionE(1280, BN)
self.Mixed7c = InceptionE(2048, BN)
self.fc = nn.Linear(2048, num_classes)
if aux_logits:
self.aux_conv = omth_blocks.conv_block(768,
filters=[128, 768],
kernel_sizes=[1, 5],
stride=[1, 1],
padding=[0, 0],
batch_norm=BN)
self.aux_conv.stdev = 0.01
self.aux_fc = nn.Linear(768, num_classes)
self.aux_fc.stddev = 0.001
def __init__(self,
channel,
rdb_number,
upscale_factor,
BN=nn.BatchNorm2d,
s_MSE=False,
filters=64,
group=1,
trellis=False,
img_bit=8):
super(RDN, self).__init__()
if s_MSE:
self.evaluator = Vgg16BN()
else:
self.evaluator = None
self.trellis = trellis
#self.group_conv1 = block.conv_block(channel, [filters * group, filters * group], kernel_sizes=[3, 3],
#stride=[1, 1], padding=[1, 1], groups=[group] * 2, name="block1", batch_norm=BN)
#self.conv1 = nn.Conv2d(filters * group, out_channels=filters, kernel_size=1, padding=0, stride=1)
rgb_mean = 0.5
rgb_std = 1.0
self.sub_mean = module.MeanShift(2**img_bit,
rgb_mean,
rgb_std,
channel=channel)
self.add_mean = module.MeanShift(2**img_bit,
rgb_mean,
rgb_std,
sign=1,
channel=1)
self.SFF1 = nn.Conv2d(in_channels=channel,
out_channels=filters,
kernel_size=3,
padding=1,
stride=1)
self.SFF2 = nn.Conv2d(in_channels=filters,
out_channels=filters,
kernel_size=3,
padding=1,
stride=1)
self.RDB1 = RDB(nb_layers=rdb_number,
input_dim=filters,
growth_rate=filters)
self.RDB2 = RDB(nb_layers=rdb_number,
input_dim=filters,
growth_rate=filters)
self.RDB3 = RDB(nb_layers=rdb_number,
input_dim=filters,
growth_rate=filters)
self.GFF1 = nn.Conv2d(in_channels=filters * 3,
out_channels=filters,
kernel_size=1,
padding=0)
self.GFF2 = nn.Conv2d(in_channels=filters,
out_channels=filters,
kernel_size=3,
padding=1)
self.upconv = nn.Conv2d(in_channels=filters,
out_channels=(filters * upscale_factor *
upscale_factor),
kernel_size=3,
padding=1)
self.pixelshuffle = nn.PixelShuffle(upscale_factor)
#self.conv2 = nn.Conv2d(in_channels=filters, out_channels=1, kernel_size=3, padding=1)
if trellis:
self.trellis = module.Trellis_Structure(filters=filters,
depth=4,
out_depth=1)
self.norm_conv1 = block.conv_block(filters, [filters, filters, 1],
kernel_sizes=[3, 3, 1],
stride=[1, 1, 1],
padding=[1, 1, 0],
groups=[1] * 3,
name="norm_conv1",
batch_norm=None,
activation=None)
"""
self.norm_conv2 = block.conv_block(64, [64, 64, 32], kernel_sizes=[3, 1, 3], stride=[1, 1, 1],
padding=[1, 0, 1], groups=[1] * 3, name="norm_conv2", batch_norm=BN,
activation=None)
self.norm_conv3 = block.conv_block(32, [32, 16, 16, 1], kernel_sizes=[1, 3, 3, 3], stride=[1, 1, 1, 1],
padding=[0, 1, 1, 1], groups=[1] * 4, name="norm_conv3", batch_norm=BN,
activation=None)
"""
self.mae = nn.L1Loss()
self.s_mse_loss = nn.MSELoss()
def __init__(self,
inchannel=3,
BN=nn.BatchNorm2d,
group=1,
s_MSE=False,
SA=True):
super(ProbaV_basic, self).__init__()
if s_MSE:
self.evaluator = Vgg16BN()
else:
self.evaluator = None
self.SA = SA
self.down_conv1 = block.conv_block(
inchannel, [48 * group, 128 * group, 128 * group],
kernel_sizes=[3, 3, 1],
stride=[2, 1, 1],
padding=[1, 1, 0],
groups=[group] * 3,
name="down_block1",
batch_norm=BN)
self.down_conv2 = block.conv_block(
128 * group, [256 * group, 256 * group, 256 * group],
kernel_sizes=[3, 3, 1],
stride=[2, 1, 1],
padding=[1, 1, 0],
groups=[group] * 3,
name="down_block2",
batch_norm=BN)
self.norm_conv1 = block.conv_block(256 * group, [256, 256, 256],
kernel_sizes=[3, 3, 3],
stride=[1] * 3,
padding=[2, 1, 1],
groups=[1] * 3,
dilation=[2, 1, 1],
name="norm_conv1",
batch_norm=BN)
if SA:
self.self_attn = Self_Attn(256)
self.up_conv1 = block.conv_block(256, [256, 256, 256],
kernel_sizes=[5, 3, 3],
stride=[3, 1, 1],
padding=[1, 1, 1],
groups=[1] * 3,
name="up_block1",
batch_norm=BN,
transpose=[True, False, False])
self.norm_conv2 = block.conv_block(256, [256, 256, 256],
kernel_sizes=[1, 3, 3],
stride=[1, 1, 1],
padding=[0, 1, 1],
groups=[1] * 3,
name="norm_conv2",
batch_norm=BN)
self.up_conv2 = block.conv_block(256, [256, 256, 256],
kernel_sizes=[4, 3, 3],
stride=[2, 1, 1],
padding=[1, 1, 1],
groups=[1] * 3,
name="up_block2",
batch_norm=BN,
transpose=[True, False, False])
self.norm_conv3 = block.conv_block(256, [256, 128, 128],
kernel_sizes=[1, 3, 3],
stride=[1, 1, 1],
padding=[0, 1, 1],
groups=[1] * 3,
name="norm_conv3",
batch_norm=BN)
self.up_conv3 = block.conv_block(128, [128, 128, 128],
kernel_sizes=[4, 3, 3],
stride=[2, 1, 1],
padding=[1, 1, 1],
groups=[1] * 3,
name="up_block3",
batch_norm=BN,
transpose=[True, False, False])
self.norm_conv4 = block.conv_block(128, [128, 64, 64],
kernel_sizes=[1, 3, 3],
stride=[1, 1, 1],
padding=[0, 1, 1],
groups=[1] * 3,
name="norm_conv4",
batch_norm=BN,
activation=None)
self.norm_conv5 = block.conv_block(64, [64, 32, 32],
kernel_sizes=[1, 3, 3],
stride=[1, 1, 1],
padding=[0, 1, 1],
groups=[1] * 3,
name="norm_conv5",
batch_norm=BN,
activation=None)
self.norm_conv6 = block.conv_block(32, [32, 16, 16, 1],
kernel_sizes=[1, 3, 3, 3],
stride=[1, 1, 1, 1],
padding=[0, 1, 1, 1],
groups=[1] * 4,
name="norm_conv6",
batch_norm=BN,
activation=None)
def create_backbone_model(self):
if self.extra_layer_filters is None:
self.extra_layer_filters = [512, 384, 384, 256, 256, 256, 256, 256]
# Prepare VGG-16 net with batch normalization
vgg16_model = vgg16_bn(pretrained=True)
net = list(vgg16_model.children())[0]
# Replace the maxout with ceil in vanilla vgg16 net
ceil_maxout = nn.MaxPool2d(kernel_size=2,
stride=2,
padding=0,
dilation=1,
ceil_mode=True)
net = [ceil_maxout if type(n) is nn.MaxPool2d else n for n in net]
# Basic VGG Layers
self.conv_module_name.append("conv_1")
self.conv_module.append(nn.Sequential(*net[:7]))
self.conv_module_name.append("conv_2")
self.conv_module.append(nn.Sequential(*net[7:14]))
self.conv_module_name.append("conv_3")
self.conv_module.append(nn.Sequential(*net[14:24]))
self.fpn_back.update({"conv_3": FPN_block(512, 256)})
self.conv_module_name.append("conv_4")
self.conv_module.append(nn.Sequential(*net[24:34]))
self.fpn_back.update({"conv_4": FPN_block(512, 512)})
self.conv_module_name.append("conv_5")
self.conv_module.append(nn.Sequential(*net[34:44]))
self.fpn_back.update({
"conv_5":
FPN_block(self.extra_layer_filters[1], 512, upscale_factor=1)
})
# Extra Layers
self.conv_module_name.append("extra_1")
self.conv_module.append(
omth_blocks.conv_block(
512,
[self.extra_layer_filters[0], self.extra_layer_filters[1]],
kernel_sizes=[3, 3],
stride=[1, 1],
padding=[2, 3],
dilation=[2, 3],
batch_norm=self.batch_norm))
self.fpn_back.update({
"extra_1":
FPN_block(self.extra_layer_filters[3],
self.extra_layer_filters[1],
upscale_factor=2)
})
self.conv_module_name.append("extra_2")
self.conv_module.append(
omth_blocks.conv_block(self.extra_layer_filters[1],
kernel_sizes=[3, 1],
filters=[
self.extra_layer_filters[2],
self.extra_layer_filters[3]
],
stride=[1, 2],
padding=[1, 0],
batch_norm=self.batch_norm))
self.fpn_back.update({
"extra_2":
FPN_block(self.extra_layer_filters[5],
self.extra_layer_filters[3],
upscale_factor=2)
})
self.conv_module_name.append("extra_3")
self.conv_module.append(
omth_blocks.conv_block(self.extra_layer_filters[3],
kernel_sizes=[3, 1],
filters=[
self.extra_layer_filters[4],
self.extra_layer_filters[5]
],
stride=[1, 2],
padding=[1, 0],
batch_norm=self.batch_norm))
self.fpn_back.update({
"extra_3":
FPN_block(self.extra_layer_filters[7],
self.extra_layer_filters[5],
upscale_factor=2)
})
self.conv_module_name.append("extra_4")
self.conv_module.append(
omth_blocks.conv_block(self.extra_layer_filters[5],
kernel_sizes=[3, 1],
filters=[
self.extra_layer_filters[6],
self.extra_layer_filters[7]
],
stride=[1, 2],
padding=[1, 0],
batch_norm=self.batch_norm))
