def __init__(self, args, classes=21, dataset='pascal'):
super().__init__()
# =============================================================
# BASE NETWORK
# =============================================================
self.base_net = EESPNet(args) #imagenet model
del self.base_net.classifier
del self.base_net.level5
del self.base_net.level5_0
config = self.base_net.config
#=============================================================
# SEGMENTATION NETWORK
#=============================================================
dec_feat_dict = {
'pascal': 16,
'city': 16,
'coco': 32,
'greenhouse': 16,
'ishihara': 16,
'camvid': 16
}
base_dec_planes = dec_feat_dict[dataset]
dec_planes = [
4 * base_dec_planes, 3 * base_dec_planes, 2 * base_dec_planes,
classes
]
pyr_plane_proj = min(classes // 2, base_dec_planes)
self.bu_dec_l1 = EfficientPyrPool(in_planes=config[3],
proj_planes=pyr_plane_proj,
out_planes=dec_planes[0])
self.bu_dec_l2 = EfficientPyrPool(in_planes=dec_planes[0],
proj_planes=pyr_plane_proj,
out_planes=dec_planes[1])
self.bu_dec_l3 = EfficientPyrPool(in_planes=dec_planes[1],
proj_planes=pyr_plane_proj,
out_planes=dec_planes[2])
self.bu_dec_l4 = EfficientPyrPool(in_planes=dec_planes[2],
proj_planes=pyr_plane_proj,
out_planes=dec_planes[3],
last_layer_br=False)
self.merge_enc_dec_l2 = EfficientPWConv(config[2], dec_planes[0])
self.merge_enc_dec_l3 = EfficientPWConv(config[1], dec_planes[1])
self.merge_enc_dec_l4 = EfficientPWConv(config[0], dec_planes[2])
self.bu_br_l2 = nn.Sequential(nn.BatchNorm2d(dec_planes[0]),
nn.PReLU(dec_planes[0]))
self.bu_br_l3 = nn.Sequential(nn.BatchNorm2d(dec_planes[1]),
nn.PReLU(dec_planes[1]))
self.bu_br_l4 = nn.Sequential(nn.BatchNorm2d(dec_planes[2]),
nn.PReLU(dec_planes[2]))
#self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
self.init_params()
def __init__(self, args, classes=21, dataset='pascal'):
super().__init__()
# =============================================================
# BASE NETWORK
# =============================================================
self.base_net = CNNModel(args) #imagenet model
del self.base_net.classifier
del self.base_net.level5
config = self.base_net.out_channel_map
#=============================================================
# SEGMENTATION NETWORK
#=============================================================
dec_feat_dict = {'pascal': 16, 'city': 16, 'coco': 32}
base_dec_planes = dec_feat_dict[dataset]
dec_planes = [
4 * base_dec_planes, 3 * base_dec_planes, 2 * base_dec_planes,
classes
]
pyr_plane_proj = min(classes // 2, base_dec_planes)
# dimensions in variable names are shown for an input of 256x256
self.eff_pool_16x16 = EfficientPyrPool(in_planes=config[3],
proj_planes=pyr_plane_proj,
out_planes=dec_planes[0])
self.eff_pool_32x32 = EfficientPyrPool(in_planes=dec_planes[0],
proj_planes=pyr_plane_proj,
out_planes=dec_planes[1])
self.eff_pool_64x64 = EfficientPyrPool(in_planes=dec_planes[1],
proj_planes=pyr_plane_proj,
out_planes=dec_planes[2])
self.eff_pool_128x128 = EfficientPyrPool(in_planes=dec_planes[2],
proj_planes=pyr_plane_proj,
out_planes=dec_planes[3],
last_layer_br=False)
self.proj_enc_32x32 = EfficientPWConv(config[2], dec_planes[0])
self.proj_enc_64x64 = EfficientPWConv(config[1], dec_planes[1])
self.proj_enc_128x128 = EfficientPWConv(config[0], dec_planes[2])
self.bu_br_32x32 = nn.Sequential(nn.BatchNorm2d(dec_planes[0]),
nn.PReLU(dec_planes[0]))
self.bu_br_64x64 = nn.Sequential(nn.BatchNorm2d(dec_planes[1]),
nn.PReLU(dec_planes[1]))
self.bu_br_128x128 = nn.Sequential(nn.BatchNorm2d(dec_planes[2]),
nn.PReLU(dec_planes[2]))
self.upsample = nn.Upsample(scale_factor=2,
mode='bilinear',
align_corners=True)
self.init_params()
def __init__(self, args, extra_layer):
'''
:param classes: number of classes in the dataset. Default is 1000 for the ImageNet dataset
:param s: factor that scales the number of output feature maps
'''
super(SSDNet512, self).__init__()
# =============================================================
# BASE NETWORK
# =============================================================
self.basenet = CNNModel(args)
# delete the classification layer
del self.basenet.classifier
# retrive the basenet configuration
base_net_config = self.basenet.out_channel_map
config = base_net_config[:4] + [base_net_config[5]]
# add configuration for SSD version
config += [1024, 512, 256, 128]
# =============================================================
# EXTRA LAYERS for DETECTION
# =============================================================
self.extra_level6 = extra_layer(config[4], config[5])
self.extra_level7 = extra_layer(config[5], config[6])
self.extra_level8 = extra_layer(config[6], config[7])
self.extra_level9 = nn.Sequential(
nn.Conv2d(config[7],
config[8],
kernel_size=3,
stride=2,
bias=False,
padding=1), nn.ReLU(inplace=True))
# =============================================================
# EXTRA LAYERS for Bottom-up decoding
# =============================================================
from nn_layers.efficient_pyramid_pool import EfficientPyrPool
in_features = config[5] + config[6]
out_features = config[5]
red_factor = 5
self.bu_4x4_8x8 = EfficientPyrPool(in_planes=in_features,
proj_planes=out_features //
red_factor,
out_planes=out_features)
in_features = config[4] + config[5]
out_features = config[4]
self.bu_8x8_16x16 = EfficientPyrPool(in_planes=in_features,
proj_planes=out_features //
red_factor,
out_planes=out_features)
in_features = config[4] + config[3]
out_features = config[3]
self.bu_16x16_32x32 = EfficientPyrPool(in_planes=in_features,
proj_planes=out_features //
red_factor,
out_planes=out_features)
in_features = config[3] + config[2]
out_features = config[2]
self.bu_32x32_64x64 = EfficientPyrPool(in_planes=in_features,
proj_planes=out_features //
red_factor,
out_planes=out_features)
self.config = config
def __init__(self, args, extra_layer):
'''
:param classes: number of classes in the dataset. Default is 1000 for the ImageNet dataset
:param s: factor that scales the number of output feature maps
'''
super(ESPNetv2SSD300, self).__init__()
# =============================================================
# BASE NETWORK
# =============================================================
self.basenet = EESPNet(args)
# delete the classification layer
del self.basenet.classifier
# retrive the basenet configuration
base_net_config = self.basenet.config
config = base_net_config[:4] + [base_net_config[5]]
# add configuration for SSD version
config += [1024, 512, 256]
# =============================================================
# EXTRA LAYERS for DETECTION
# =============================================================
self.extra_level6 = extra_layer(config[4], config[5])
self.extra_level7 = extra_layer(config[5], config[6])
self.extra_level8 = nn.Sequential(
nn.Conv2d(config[6],
config[6],
kernel_size=3,
stride=2,
bias=False,
padding=1), nn.BatchNorm2d(config[6]),
nn.ReLU(inplace=True),
nn.Conv2d(config[6],
config[7],
kernel_size=2,
stride=2,
bias=False), nn.ReLU(inplace=True))
# =============================================================
# EXTRA LAYERS for Bottom-up decoding
# =============================================================
from nn_layers.efficient_pyramid_pool import EfficientPyrPool
in_features = config[5] + config[6]
out_features = config[5]
red_factor = 5
self.bu_3x3_5x5 = EfficientPyrPool(in_planes=in_features,
proj_planes=out_features //
red_factor,
out_planes=out_features)
in_features = config[4] + config[5]
out_features = config[4]
self.bu_5x5_10x10 = EfficientPyrPool(in_planes=in_features,
proj_planes=out_features //
red_factor,
out_planes=out_features)
in_features = config[4] + config[3]
out_features = config[3]
self.bu_10x10_19x19 = EfficientPyrPool(in_planes=in_features,
proj_planes=out_features //
red_factor,
out_planes=out_features)
in_features = config[3] + config[2]
out_features = config[2]
self.bu_19x19_38x38 = EfficientPyrPool(in_planes=in_features,
proj_planes=out_features //
red_factor,
out_planes=out_features)
self.config = config
def __init__(self, args, classes=21, dataset='pascal', dense_fuse=False, trainable_fusion=True):
super().__init__()
# =============================================================
# BASE NETWORK
# =============================================================
#
# RGB
#
self.base_net = EESPNet(args) #imagenet model
del self.base_net.classifier
del self.base_net.level5
del self.base_net.level5_0
config = self.base_net.config
#
# Depth
#
tmp_args = copy.deepcopy(args)
tmp_args.channels = 1
self.depth_base_net = EESPNet(tmp_args)
del self.depth_base_net.classifier
del self.depth_base_net.level5
del self.depth_base_net.level5_0
self.fusion_gate_level1 = FusionGate(nchannel=32, is_trainable=trainable_fusion)
self.fusion_gate_level2 = FusionGate(nchannel=128, is_trainable=trainable_fusion)
self.fusion_gate_level3 = FusionGate(nchannel=256, is_trainable=trainable_fusion)
self.fusion_gate_level4 = FusionGate(nchannel=512, is_trainable=trainable_fusion)
# Layer 1
# self.depth_encoder_level1 = nn.Sequential(
# CBR(nIn=1, nOut=32, kSize=3, stride=2), # Input: 3, Ouput: 16, kernel: 3
# CBR(nIn=32, nOut=32, kSize=3), # Input: 3, Ouput: 16, kernel: 3
# )
#
# # Level 2
# self.depth_encoder_level2 = nn.Sequential(
## C(nIn=32, nOut=128, kSize=1), # Pixel-wise conv
## CBR(nIn=128, nOut=128, kSize=3, stride=2, groups=128) # Depth-wise conv
# CBR(nIn=32, nOut=128, kSize=3, stride=2), # Downsample
# CBR(nIn=128, nOut=128, kSize=3),
# CBR(nIn=128, nOut=128, kSize=3)
# )
#
# # Level 3
# self.depth_encoder_level3 = nn.Sequential(
## C(nIn=128, nOut=256, kSize=1), # Pixel-wise conv
## CBR(nIn=256, nOut=256, kSize=3, groups=256) # Depth-wise conv
# CBR(nIn=128, nOut=256, kSize=3, stride=2),
# CBR(nIn=256, nOut=256, kSize=3),
# CBR(nIn=256, nOut=256, kSize=3),
# CBR(nIn=256, nOut=256, kSize=3)
#
# )
#
# # Level 4
# self.depth_encoder_level4 = nn.Sequential(
# CBR(nIn=256, nOut=512, kSize=3, stride=2), # Pixel-wise conv
# CBR(nIn=512, nOut=512, kSize=3),
# CBR(nIn=512, nOut=512, kSize=3),
# CBR(nIn=512, nOut=512, kSize=3)
# )
# 112 L1
#=============================================================
# SEGMENTATION NETWORK
#=============================================================
dec_feat_dict={
'pascal': 16,
'city': 16,
'coco': 32,
'greenhouse': 16,
'ishihara': 16,
'sun': 16,
'camvid': 16
}
base_dec_planes = dec_feat_dict[dataset]
dec_planes = [4*base_dec_planes, 3*base_dec_planes, 2*base_dec_planes, classes]
pyr_plane_proj = min(classes //2, base_dec_planes)
self.bu_dec_l1 = EfficientPyrPool(in_planes=config[3], proj_planes=pyr_plane_proj,
out_planes=dec_planes[0])
self.bu_dec_l2 = EfficientPyrPool(in_planes=dec_planes[0], proj_planes=pyr_plane_proj,
out_planes=dec_planes[1])
self.bu_dec_l3 = EfficientPyrPool(in_planes=dec_planes[1], proj_planes=pyr_plane_proj,
out_planes=dec_planes[2])
self.bu_dec_l4 = EfficientPyrPool(in_planes=dec_planes[2], proj_planes=pyr_plane_proj,
out_planes=dec_planes[3], last_layer_br=False)
self.merge_enc_dec_l2 = EfficientPWConv(config[2], dec_planes[0])
self.merge_enc_dec_l3 = EfficientPWConv(config[1], dec_planes[1])
self.merge_enc_dec_l4 = EfficientPWConv(config[0], dec_planes[2])
self.bu_br_l2 = nn.Sequential(nn.BatchNorm2d(dec_planes[0]),
nn.PReLU(dec_planes[0])
)
self.bu_br_l3 = nn.Sequential(nn.BatchNorm2d(dec_planes[1]),
nn.PReLU(dec_planes[1])
)
self.bu_br_l4 = nn.Sequential(nn.BatchNorm2d(dec_planes[2]),
nn.PReLU(dec_planes[2])
)
#self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
self.init_params()
self.dense_fuse = dense_fuse
def __init__(self,
args,
classes=21,
dataset='pascal',
dense_fuse=False,
trainable_fusion=True,
aux_layer=2,
fix_pyr_plane_proj=False,
in_channels=32,
spatial=True):
super().__init__()
# =============================================================
# BASE NETWORK
# =============================================================
#
# RGB
#
self.base_net = EESPNet(args) #imagenet model
del self.base_net.classifier
del self.base_net.level5
del self.base_net.level5_0
config = self.base_net.config
#
# Depth
#
tmp_args = copy.deepcopy(args)
tmp_args.channels = 1
self.depth_base_net = EESPNet(tmp_args)
del self.depth_base_net.classifier
del self.depth_base_net.level5
del self.depth_base_net.level5_0
self.fusion_gate_level1 = FusionGate(nchannel=32,
is_trainable=trainable_fusion)
self.fusion_gate_level2 = FusionGate(nchannel=128,
is_trainable=trainable_fusion)
self.fusion_gate_level3 = FusionGate(nchannel=256,
is_trainable=trainable_fusion)
self.fusion_gate_level4 = FusionGate(nchannel=512,
is_trainable=trainable_fusion)
#=============================================================
# SEGMENTATION NETWORK
#=============================================================
dec_feat_dict = {
'pascal': 16,
'city': 16,
'coco': 32,
'greenhouse': 16,
'ishihara': 16,
'sun': 16,
'camvid': 16,
'forest': 16
}
base_dec_planes = dec_feat_dict[dataset]
dec_planes = [
4 * base_dec_planes, 3 * base_dec_planes, 2 * base_dec_planes,
classes
]
if fix_pyr_plane_proj:
pyr_plane_proj = base_dec_planes
else:
pyr_plane_proj = min(classes // 2, base_dec_planes)
self.bu_dec_l1 = EfficientPyrPool(in_planes=config[3],
proj_planes=pyr_plane_proj,
out_planes=dec_planes[0])
self.bu_dec_l2 = EfficientPyrPool(in_planes=dec_planes[0],
proj_planes=pyr_plane_proj,
out_planes=dec_planes[1])
self.bu_dec_l3 = EfficientPyrPool(in_planes=dec_planes[1],
proj_planes=pyr_plane_proj,
out_planes=dec_planes[2])
self.bu_dec_l4 = EfficientPyrPool(in_planes=dec_planes[2],
proj_planes=pyr_plane_proj,
out_planes=dec_planes[3],
last_layer_br=False)
self.merge_enc_dec_l2 = EfficientPWConv(config[2], dec_planes[0])
self.merge_enc_dec_l3 = EfficientPWConv(config[1], dec_planes[1])
self.merge_enc_dec_l4 = EfficientPWConv(config[0], dec_planes[2])
self.bu_br_l2 = nn.Sequential(nn.BatchNorm2d(dec_planes[0]),
nn.PReLU(dec_planes[0]))
self.bu_br_l3 = nn.Sequential(nn.BatchNorm2d(dec_planes[1]),
nn.PReLU(dec_planes[1]))
self.bu_br_l4 = nn.Sequential(nn.BatchNorm2d(dec_planes[2]),
nn.PReLU(dec_planes[2]))
# Auxiliary branch
self.aux_layer = aux_layer
if aux_layer >= 0 and aux_layer < 3:
self.aux_decoder = EfficientPyrPool(
in_planes=dec_planes[aux_layer],
proj_planes=pyr_plane_proj,
out_planes=dec_planes[3],
last_layer_br=False)
#self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
self.init_params()
self.dense_fuse = dense_fuse
#
# Traversability module
#
self.trav_module = LabelProbEstimator(in_channels=in_channels,
spatial=spatial)
# Register
self.activation = {}
def get_activation(name):
def hook(model, input, output):
self.activation[name] = output.detach()
return hook
self.bu_dec_l4.merge_layer[2].register_forward_hook(
get_activation('output_main'))
self.aux_decoder.merge_layer[2].register_forward_hook(
get_activation('output_aux'))