def test_feature_extraction_methods_equivalence(self):
model = models.resnet18(**self.model_defaults).eval()
return_layers = {
'layer1': 'layer1',
'layer2': 'layer2',
'layer3': 'layer3',
'layer4': 'layer4'
}
ilg_model = IntermediateLayerGetter(model, return_layers).eval()
fx_model = self._create_feature_extractor(model, return_layers)
# Check that we have same parameters
for (n1, p1), (n2, p2) in zip(ilg_model.named_parameters(),
fx_model.named_parameters()):
assert n1 == n2
assert p1.equal(p2)
# And that ouputs match
with torch.no_grad():
ilg_out = ilg_model(self.inp)
fgn_out = fx_model(self.inp)
assert all(k1 == k2 for k1, k2 in zip(ilg_out.keys(), fgn_out.keys()))
for k in ilg_out.keys():
assert ilg_out[k].equal(fgn_out[k])
class Segmentator(nn.Module):
def __init__(self,
num_classes,
encoder,
img_size=(512, 512),
shallow_decoder=False):
super().__init__()
self.low_feat = IntermediateLayerGetter(encoder,
{"layer1": "layer1"}) #.cuda()
self.encoder = IntermediateLayerGetter(encoder,
{"layer4": "out"}) #.cuda()
# n_classes, encoder_dim, img_size, low_level_dim, rates
self.decoder = Decoder(num_classes,
512,
img_size,
low_level_dim=64,
rates=[1, 6, 12, 18])
self.num_classes = num_classes
def forward(self, x):
self.low_feat.eval()
self.encoder.eval()
with torch.no_grad():
# This is possible since gradients are not being updated
low_level_feat = self.low_feat(x)['layer1']
enc_feat = self.encoder(x)['out']
segmentation = self.decoder(enc_feat, low_level_feat)
if self.num_classes == 1:
segmentation = torch.sigmoid(segmentation)
return segmentation
def __init__(self, num_convs=3, fine_tune=False):
super(MyModel, self).__init__()
assert 8 >= num_convs > 1, "Cannot have less than 1 or greater than 8 convolutional+pooling layers."
self.num_convs = num_convs
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
# Simple Segmentation Model
self.backbone = models.resnet50(
pretrained=True, replace_stride_with_dilation=[False, True, True])
self.backbone = IntermediateLayerGetter(
self.backbone, return_layers={'layer4': 'out'})
if not fine_tune:
for p in self.backbone.parameters():
p.requires_grad = False
in_channels = 1028
classes = 200 # actually classes
# Graph convolution layers and graph pooling layers
self.classifier = FCNHead(in_channels, classes)
self.convolutions = dict()
for conv in range(num_convs):
self.convolutions['conv' + str(conv + 1)] = GCNConv(2048, 2048).to(
self.device)
self.convolutions['pool' + str(conv + 1)] = TopKPooling(
2048, ratio=.6).to(self.device)
#Final Output
self.lin1 = torch.nn.Linear(2048, 1024).to(self.device)
self.lin2 = torch.nn.Linear(1024, 1024).to(self.device)
self.lin3 = torch.nn.Linear(1024, classes).to(self.device)
self.act1 = torch.nn.ReLU().to(self.device)
self.act2 = torch.nn.ReLU().to(self.device)
def __init__(self, train_backbone: bool, channel_768: bool):
super().__init__()
if tiny:
self.body = IntermediateLayerGetter(timm.create_model(
'vit_small_patch16_224', pretrained=True),
return_layers={'blocks': '0'})
self.tiny = True
else:
self.body = IntermediateLayerGetter(timm.create_model(
'vit_base_patch16_384', pretrained=True),
return_layers={'blocks': '0'})
self.tiny = False
if train_backbone:
for name, parameter in self.body.named_parameters():
parameter.requires_grad_(True)
else:
for name, parameter in self.body.named_parameters():
parameter.requires_grad_(False)
if channel_768:
self.num_channels = [768]
else:
self.num_channels = [2048]
self.strides = [32]
def __init__(self,
num_classes,
encoder,
att_type='additive',
img_size=(512, 512)):
super().__init__()
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
self.low_feat = IntermediateLayerGetter(encoder, {
"layer1": "layer1"
}).to(self.device)
self.encoder = IntermediateLayerGetter(encoder, {
"layer4": "out"
}).to(self.device)
# For resnet18
encoder_dim = 512
low_level_dim = 64
self.num_classes = num_classes
self.class_encoder = nn.Linear(num_classes, 512)
self.attention_enc = Attention(encoder_dim, att_type)
self.decoder = Decoder(2,
encoder_dim,
img_size,
low_level_dim=low_level_dim,
rates=[1, 6, 12, 18])
def __init__(self,
arch='resnet50',
pretrained=True,
input_size=[640, 640]):
super().__init__()
if 'resnet' in arch:
from torchvision.models import resnet
self.body = resnet.__dict__[arch](pretrained)
return_layers = {'layer2': 0, 'layer3': 1, 'layer4': 2}
self.body = IntermediateLayerGetter(self.body, return_layers)
channels = [512, 1024, 2048]
else: # todo: retina face with mobilenet as body
from torchvision.models import mobilenet_v2
self.body = mobilenet_v2(pretrained=True)
return_layers = {}
self.body = IntermediateLayerGetter(self.body, return_layers)
channels = []
self.fpn = FPN(channels, 256)
self.classifications = [
ClassHead(512, 3) for _ in range(len(return_layers))
]
self.boundingboxes = [
BboxHead(512, 3) for _ in range(len(return_layers))
]
self.landmarks = [LarkHead(512, 3) for _ in range(len(return_layers))]
self.anchors = Anchors()(input_size=input_size)
self.pts_tfms = RegressionTransform()
def __init__(self):
super(MyGCNNet, self).__init__()
# BASE
backbone = resnet.__dict__["resnet50"](pretrained=True, replace_stride_with_dilation=[False, False, True])
return_layers = {'relu': 'e0', 'layer1': 'e1', 'layer2': 'e2', 'layer3': 'e3', 'layer4': 'e4'}
self.backbone = IntermediateLayerGetter(backbone, return_layers=return_layers)
for para in self.backbone.named_parameters():
if "bn" in para[0]:
para[1].requires_grad = False
pass
# Convert
self.relu = nn.ReLU(inplace=True)
self.convert5 = nn.Conv2d(2048, 512, 1, 1, bias=False) # 25
self.convert4 = nn.Conv2d(1024, 512, 1, 1, bias=False) # 25
self.convert3 = nn.Conv2d(512, 256, 1, 1, bias=False) # 50
self.convert2 = nn.Conv2d(256, 256, 1, 1, bias=False) # 100
self.convert1 = nn.Conv2d(64, 128, 1, 1, bias=False) # 200
# DEEP POOL
deep_pool = [[512, 512, 256, 256, 128], [512, 256, 256, 128, 128]]
self.deep_pool5 = DeepPoolLayer(deep_pool[0][0], deep_pool[1][0], True, True)
self.deep_pool4 = DeepPoolLayer(deep_pool[0][1], deep_pool[1][1], True, True)
self.deep_pool3 = DeepPoolLayer(deep_pool[0][2], deep_pool[1][2], True, True)
self.deep_pool2 = DeepPoolLayer(deep_pool[0][3], deep_pool[1][3], True, True)
self.deep_pool1 = DeepPoolLayer(deep_pool[0][4], deep_pool[1][4], False, False)
# ScoreLayer
score = 128
self.score = nn.Conv2d(score, 1, 1, 1)
pass
def backbone(backbone_name, dataset):
if 'BSDS' in dataset:
dilation = [False, True, True]
elif 'NYUD' in dataset:
dilation = [False, False, True]
return_layers = {
'relu': 'feat1',
'layer1': 'feat2',
'layer2': 'feat3',
'layer3': 'feat4',
'layer4': 'feat5'
}
if backbone_name == 'vgg':
backbone = VGG()
elif backbone_name == 'resnet50':
backbone = resnet50(pretrained=False,
replace_stride_with_dilation=dilation)
state_dict = load_state_dict_from_url(model_urls['resnet50'],
progress=True)
state_dict = rmBN_in_ResNet(state_dict)
backbone.load_state_dict(state_dict)
backbone = IntermediateLayerGetter(backbone,
return_layers=return_layers)
elif backbone_name == 'resnet101':
backbone = resnet101(pretrained=False,
replace_stride_with_dilation=dilation)
state_dict = load_state_dict_from_url(model_urls['resnet101'],
progress=True)
state_dict = rmBN_in_ResNet(state_dict)
backbone.load_state_dict(state_dict)
backbone = IntermediateLayerGetter(backbone,
return_layers=return_layers)
elif backbone_name == 'resnext50':
backbone = resnext50_32x4d(pretrained=False,
replace_stride_with_dilation=dilation)
state_dict = load_state_dict_from_url(model_urls['resnext50_32x4d'],
progress=True)
state_dict = rmBN_in_ResNet(state_dict)
backbone.load_state_dict(state_dict)
backbone = IntermediateLayerGetter(backbone,
return_layers=return_layers)
elif backbone_name == 'resnext101':
backbone = resnext101_32x8d(pretrained=False,
replace_stride_with_dilation=dilation)
state_dict = load_state_dict_from_url(model_urls['resnext101_32x8d'],
progress=True)
state_dict = rmBN_in_ResNet(state_dict)
backbone.load_state_dict(state_dict)
backbone = IntermediateLayerGetter(backbone,
return_layers=return_layers)
return backbone
def __init__(self, backbone_name:str="resnet50", train_layers=['layer2', 'layer3', 'layer4']):
super().__init__()
backbone = getattr(torchvision.models, backbone_name)(pretrained=True)
backbone.requires_grad_(False)
for name, parameter in backbone.named_parameters():
for layer in train_layers:
if layer in name:
parameter.requires_grad_(True)
# return_layers = {"layer1": "layer1", "layer2": "layer2", "layer3": "layer3", "layer4": "layer4"}
return_layers = {"layer4": "layer4"}
self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)
for k, v in self.body.items():
setattr(self, k, v)
self.num_channels = 512 if backbone_name in ('resnet18', 'resnet34') else 2048
class AttSegmentator(nn.Module):
def __init__(self,
num_classes,
encoder,
att_type='additive',
img_size=(512, 512)):
super().__init__()
self.low_feat = IntermediateLayerGetter(encoder, {
"layer1": "layer1"
}).cuda()
self.encoder = IntermediateLayerGetter(encoder, {
"layer4": "out"
}).cuda()
# For resnet18
encoder_dim = 512
low_level_dim = 64
self.num_classes = num_classes
self.class_encoder = nn.Linear(num_classes, 512)
self.attention_enc = Attention(encoder_dim, att_type)
self.decoder = Decoder(2,
encoder_dim,
img_size,
low_level_dim=low_level_dim,
rates=[1, 6, 12, 18])
def forward(self, x, v_class, out_att=False):
self.low_feat.eval()
self.encoder.eval()
with torch.no_grad():
low_level_feat = self.low_feat(x)['layer1']
enc_feat = self.encoder(x)['out']
query = self.class_encoder(v_class)
shape = enc_feat.shape
enc_feat = enc_feat.permute(0, 2, 3, 1).contiguous().view(
shape[0], -1, shape[1])
x_enc, attention = self.attention_enc(enc_feat, query)
x_enc = x_enc.view(shape)
segmentation = self.decoder(x_enc, low_level_feat)
if out_att:
return segmentation, attention
return segmentation
def __init__(self, network_name='resnet50', half=False, phase='test'):
super(RetinaFace, self).__init__()
self.half_inference = half
cfg = generate_config(network_name)
self.backbone = cfg['name']
self.model_name = f'retinaface_{network_name}'
self.cfg = cfg
self.phase = phase
self.target_size, self.max_size = 1600, 2150
self.resize, self.scale, self.scale1 = 1., None, None
self.mean_tensor = torch.tensor([[[[104.]], [[117.]],
[[123.]]]]).to(device)
self.reference = get_reference_facial_points(default_square=True)
# Build network.
backbone = None
if cfg['name'] == 'mobilenet0.25':
backbone = MobileNetV1()
self.body = IntermediateLayerGetter(backbone, cfg['return_layers'])
elif cfg['name'] == 'Resnet50':
import torchvision.models as models
backbone = models.resnet50(pretrained=False)
self.body = IntermediateLayerGetter(backbone, cfg['return_layers'])
in_channels_stage2 = cfg['in_channel']
in_channels_list = [
in_channels_stage2 * 2,
in_channels_stage2 * 4,
in_channels_stage2 * 8,
]
out_channels = cfg['out_channel']
self.fpn = FPN(in_channels_list, out_channels)
self.ssh1 = SSH(out_channels, out_channels)
self.ssh2 = SSH(out_channels, out_channels)
self.ssh3 = SSH(out_channels, out_channels)
self.ClassHead = make_class_head(fpn_num=3,
inchannels=cfg['out_channel'])
self.BboxHead = make_bbox_head(fpn_num=3,
inchannels=cfg['out_channel'])
self.LandmarkHead = make_landmark_head(fpn_num=3,
inchannels=cfg['out_channel'])
self.to(device)
self.eval()
if self.half_inference:
self.half()
def _segm_resnet(name,
backbone_name,
num_classes,
aux,
pretrained_backbone=True):
backbone = resnet.__dict__[backbone_name](
pretrained=pretrained_backbone,
replace_stride_with_dilation=[False, True, True])
backbone.conv1 = nn.Conv2d(3,
64,
kernel_size=(7, 7),
stride=(2, 2),
padding=(3, 3),
bias=False)
return_layers = {'layer4': 'out'}
if aux:
return_layers['layer3'] = 'aux'
backbone = IntermediateLayerGetter(backbone, return_layers=return_layers)
aux_classifier = None
if aux:
inplanes = 1024
aux_classifier = FCNHead(inplanes, num_classes)
model_map = {
'fcn': (FCNHead, FCN),
}
inplanes = 2048
classifier = model_map[name][0](inplanes, num_classes)
base_model = model_map[name][1]
model = base_model(backbone, classifier, aux_classifier)
return model
def __init__(self, backbone: nn.Module, train_backbone: bool,
num_channels: int, return_interm_layers: bool,
backbone_name: str):
super().__init__()
for name, parameter in backbone.named_parameters():
if not train_backbone or ('conv5' not in name
and 'layer2' not in name
and 'layer3' not in name
and 'layer4' not in name):
parameter.requires_grad_(False)
if return_interm_layers:
if 'shufflenet' in backbone_name:
return_layers = {'conv5': "0"}
else:
return_layers = {
"layer1": "0",
"layer2": "1",
"layer3": "2",
"layer4": "3"
}
else:
if 'shufflenet' in backbone_name:
return_layers = {'conv5': "0"}
else:
return_layers = {'layer4': "0"}
self.body = IntermediateLayerGetter(backbone,
return_layers=return_layers)
self.num_channels = num_channels
def _segm_resnet(name, backbone_name, num_classes, aux, pretrained_backbone=True, **kwargs):
backbone = resnet.__dict__[backbone_name](
pretrained=pretrained_backbone,
replace_stride_with_dilation=[False, True, True],
resnet_local_path=kwargs["resnet_local_path"] if "resnet_local_path" in kwargs else None
)
return_layers = {'layer4': 'out'}
if aux:
return_layers['layer3'] = 'aux'
backbone = IntermediateLayerGetter(backbone, return_layers=return_layers)
aux_classifier = None
if aux:
inplanes = 1024
aux_classifier = FCNHead(inplanes, num_classes)
model_map = {
'deeplabv3': (DeepLabHead, DeepLabV3),
'fcn': (FCNHead, FCN),
}
inplanes = 2048
classifier = model_map[name][0](inplanes, num_classes)
base_model = model_map[name][1]
model = base_model(backbone, classifier, aux_classifier)
return model
def __init__(self, backbone_name: str, backbone: nn.Module,
train_backbone: bool, return_layers: List):
super().__init__()
for name, parameter in backbone.named_parameters():
if not train_backbone or 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:
parameter.requires_grad_(False)
self.return_layers = return_layers
return_layer_map = {}
for idx, layer in enumerate(range(2, int(return_layers[-1][-1]) + 1)):
return_layer_map['layer' + str(layer)] = str(layer)
self.body = IntermediateLayerGetter(backbone,
return_layers=return_layer_map)
self.name = backbone_name
self.num_channels_list = []
for layer in return_layers:
if layer == "layer2":
if backbone_name == "resnet18":
self.num_channels_list.append(128)
if backbone_name == "resnet50":
self.num_channels_list.append(512)
elif layer == "layer3":
if backbone_name == "resnet18":
self.num_channels_list.append(256)
if backbone_name == "resnet50":
self.num_channels_list.append(1024)
elif layer == "layer4":
if backbone_name == "resnet18":
self.num_channels_list.append(512)
if backbone_name == "resnet50":
self.num_channels_list.append(2048)
def __init__(self,
backbone_name,
pretrained_weights=None,
aux=False,
pretrained_backbone=True,
freeze_bn=False):
super(vgg_feature_extractor, self).__init__()
backbone = vgg.__dict__[backbone_name](
pretrained=pretrained_backbone,
pretrained_weights=pretrained_weights)
features, _ = list(backbone.features.children()), list(
backbone.classifier.children())
#remove pool4/pool5
features = nn.Sequential(*(features[i] for i in list(range(23)) +
list(range(24, 30))))
for i in [23, 25, 27]:
features[i].dilation = (2, 2)
features[i].padding = (2, 2)
fc6 = nn.Conv2d(512, 1024, kernel_size=3, padding=4, dilation=4)
fc7 = nn.Conv2d(1024, 1024, kernel_size=3, padding=4, dilation=4)
backbone = nn.Sequential(
*([features[i] for i in range(len(features))] +
[fc6, nn.ReLU(inplace=True), fc7,
nn.ReLU(inplace=True)]))
return_layers = {'4': 'low_fea', '32': 'out'}
self.backbone = IntermediateLayerGetter(backbone,
return_layers=return_layers)
def __init__(
self,
backbone: nn.Module,
train_backbone: bool,
num_channels: int,
return_interm_layers: bool,
):
super().__init__()
for name, parameter in backbone.named_parameters():
if (not train_backbone or "layer2" not in name
and "layer3" not in name and "layer4" not in name):
parameter.requires_grad_(False)
if return_interm_layers:
return_layers = {
"layer1": "0",
"layer2": "1",
"layer3": "2",
"layer4": "3"
}
else:
return_layers = {"layer4": 0}
self.body = IntermediateLayerGetter(backbone,
return_layers=return_layers)
print("backbone body")
print(self.body)
self.num_channels = num_channels
def segm_resnet(name, backbone_name, num_classes, output_stride, pretrained_backbone):
if output_stride == 8:
replace_stride_with_dilation=[False, True, True]
aspp_dilate = [12,24,36]
else:
replace_stride_with_dilation = [False, False, True]
aspp_dilate = [6,12,18]
backbone = resnet.__dict__[backbone_name](
pretrained=pretrained_backbone,
replace_stride_with_dilation=replace_stride_with_dilation
)
inplanes=2048
low_level_planes=256
if name == 'deeplabv3plus':
return_layers = {'layer4':'out', 'layer1': 'low_level'}
classifier = DeepLabHeadV3Plus(inplanes, low_level_planes, num_classes, aspp_dilate)
elif name=='deeplabv3':
return_layers = {'layer4': 'out'}
classifier = DeepLabHead(inplanes , num_classes , aspp_dilate)
backbone = IntermediateLayerGetter(backbone, return_layers=return_layers)
model = Deeplabv3(backbone, classifier)
return model
def resnet_backbone(backbone_arch: str, first_trainable_stage: int) -> Module:
"""return resnet backbone."""
if backbone_arch.startswith("resnext"):
# TODO - figure out why FrozenBatchNorm2d causes issues
norm_layer = BatchNorm2d
else:
norm_layer = FrozenBatchNorm2d
backbone = resnet.__dict__[backbone_arch](
pretrained=True,
norm_layer=norm_layer,
replace_stride_with_dilation=(False, False, 2),
)
backbone.eval()
# freeze layers
for name, parameter in backbone.named_parameters():
match = re.search(r"layer(\d)", name)
if not (match and int(match.group(1)) >= first_trainable_stage):
parameter.requires_grad_(False)
return_layers = {"layer2": "c3", "layer3": "c4", "layer4": "c5"}
normalizer = Normalizer()
backbone = IntermediateLayerGetter(backbone, return_layers)
backbone = Sequential(normalizer, backbone)
return backbone
def _dbnet(
arch: str,
pretrained: bool,
backbone_fn: Callable[[bool], nn.Module],
fpn_layers: List[str],
backbone_submodule: Optional[str] = None,
pretrained_backbone: bool = True,
**kwargs: Any,
) -> DBNet:
# Starting with Imagenet pretrained params introduces some NaNs in layer3 & layer4 of resnet50
pretrained_backbone = pretrained_backbone and not arch.split(
'_')[1].startswith('resnet')
pretrained_backbone = pretrained_backbone and not pretrained
# Feature extractor
backbone = backbone_fn(pretrained_backbone)
if isinstance(backbone_submodule, str):
backbone = getattr(backbone, backbone_submodule)
feat_extractor = IntermediateLayerGetter(
backbone,
{layer_name: str(idx)
for idx, layer_name in enumerate(fpn_layers)},
)
# Build the model
model = DBNet(feat_extractor, cfg=default_cfgs[arch], **kwargs)
# Load pretrained parameters
if pretrained:
load_pretrained_params(model, default_cfgs[arch]['url'])
return model
class MyGCNNet(nn.Module):
def __init__(self):
super(MyGCNNet, self).__init__()
# BASE
backbone = resnet.__dict__["resnet50"](pretrained=True, replace_stride_with_dilation=[False, False, True])
return_layers = {'relu': 'e0', 'layer1': 'e1', 'layer2': 'e2', 'layer3': 'e3', 'layer4': 'e4'}
self.backbone = IntermediateLayerGetter(backbone, return_layers=return_layers)
for para in self.backbone.named_parameters():
if "bn" in para[0]:
para[1].requires_grad = False
pass
# Convert
self.relu = nn.ReLU(inplace=True)
self.convert5 = nn.Conv2d(2048, 512, 1, 1, bias=False) # 25
self.convert4 = nn.Conv2d(1024, 512, 1, 1, bias=False) # 25
self.convert3 = nn.Conv2d(512, 256, 1, 1, bias=False) # 50
self.convert2 = nn.Conv2d(256, 256, 1, 1, bias=False) # 100
self.convert1 = nn.Conv2d(64, 128, 1, 1, bias=False) # 200
# DEEP POOL
deep_pool = [[512, 512, 256, 256, 128], [512, 256, 256, 128, 128]]
self.deep_pool5 = DeepPoolLayer(deep_pool[0][0], deep_pool[1][0], True, True)
self.deep_pool4 = DeepPoolLayer(deep_pool[0][1], deep_pool[1][1], True, True)
self.deep_pool3 = DeepPoolLayer(deep_pool[0][2], deep_pool[1][2], True, True)
self.deep_pool2 = DeepPoolLayer(deep_pool[0][3], deep_pool[1][3], True, True)
self.deep_pool1 = DeepPoolLayer(deep_pool[0][4], deep_pool[1][4], False, False)
# ScoreLayer
score = 128
self.score = nn.Conv2d(score, 1, 1, 1)
pass
def forward(self, x):
# BASE
feature = self.backbone(x)
feature1 = self.relu(self.convert1(feature["e0"])) # 128, 200
feature2 = self.relu(self.convert2(feature["e1"])) # 256, 100
feature3 = self.relu(self.convert3(feature["e2"])) # 256, 50
feature4 = self.relu(self.convert4(feature["e3"])) # 512, 25
feature5 = self.relu(self.convert5(feature["e4"])) # 512, 25
# SIZE
x_size = x.size()[2:]
merge = self.deep_pool5(feature5, feature4) # A + F
merge = self.deep_pool4(merge, feature3) # A + F
merge = self.deep_pool3(merge, feature2) # A + F
merge = self.deep_pool2(merge, feature1) # A + F
merge = self.deep_pool1(merge) # A
# ScoreLayer
merge = self.score(merge)
if x_size is not None:
merge = F.interpolate(merge, x_size, mode='bilinear', align_corners=True)
return merge, torch.sigmoid(merge)
pass
def __init__(self, latent_size=20):
super().__init__()
backbone = resnet50(False,
replace_stride_with_dilation=[False, True, True])
backbone = IntermediateLayerGetter(backbone, {"layer4": "out"})
classifier = nn.Sequential(deeplabv3.ASPP(2048, [12, 24, 36]))
self.backbone = deeplabv3.DeepLabV3(backbone, classifier, None)
self.backbone.load_state_dict(load_state_dict_from_url(
'https://download.pytorch.org/models/deeplabv3_resnet50_coco-cd0a2569.pth',
progress=True),
strict=False)
self.backbone.backbone.conv1.in_channels = 1
self.backbone.backbone.conv1.weight.data = self.backbone.backbone.conv1.weight.data.mean(
1, keepdims=True)
self.decode = nn.Sequential(
nn.Conv2d(latent_size, 256, 1),
nn.ReLU(),
nn.Conv2d(256, 2048, 1),
nn.ReLU(),
)
self.out = nn.Sequential(
nn.Conv2d(256, 256, 3, 1, 1), nn.ReLU(),
nn.Upsample(scale_factor=2), nn.Conv2d(256, 2 * 2, 3, 1, 1),
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False))
def _linknet(
arch: str,
pretrained: bool,
backbone_fn: Callable[[bool], nn.Module],
fpn_layers: List[str],
pretrained_backbone: bool = False,
**kwargs: Any,
) -> LinkNet:
pretrained_backbone = pretrained_backbone and not pretrained
# Build the feature extractor
backbone = backbone_fn(pretrained_backbone)
feat_extractor = IntermediateLayerGetter(
backbone,
{layer_name: str(idx)
for idx, layer_name in enumerate(fpn_layers)},
)
# Build the model
model = LinkNet(feat_extractor, cfg=default_cfgs[arch], **kwargs)
# Load pretrained parameters
if pretrained:
load_pretrained_params(model, default_cfgs[arch]["url"])
return model
def _segm_resnet(name,
backbone_name,
num_classes,
aux,
pretrained_backbone=False):
backbone = resnet.__dict__[backbone_name](
pretrained=pretrained_backbone,
replace_stride_with_dilation=[False, True, True])
return_layers = {"layer4": "out"}
if aux:
return_layers["layer3"] = "aux"
backbone = IntermediateLayerGetter(backbone, return_layers=return_layers)
aux_classifier = None
if aux:
inplanes = 1024
aux_classifier = FCNHead(inplanes, num_classes)
model_map = {
"deeplabv3": (DeepLabHead, DeepLabV3),
"fcn": (FCNHead, FCN),
}
inplanes = 2048
classifier = model_map[name][0](inplanes, num_classes)
base_model = model_map[name][1]
model = base_model(backbone, classifier, aux_classifier)
return model
def _segm_resnet(name, backbone_name, num_classes, aux, **kwargs):
# FIXME: 1000 and _
if isinstance(backbone_name, dict):
backbone = model_utils.which_architecture(backbone_name['arch'],
backbone_name['customs'])
elif os.path.isfile(backbone_name):
backbone, _ = model_utils.which_network_classification(
backbone_name, 1000, **kwargs)
else:
backbone = model_utils.which_architecture(backbone_name, **kwargs)
return_layers = {'layer4': 'out'}
if aux:
return_layers['layer3'] = 'aux'
backbone = IntermediateLayerGetter(backbone, return_layers=return_layers)
aux_classifier = None
if aux:
layer3 = list(backbone.layer3)[-1]
inplanes = list(layer3.children())[-2].num_features
aux_classifier = FCNHead(inplanes, num_classes)
model_map = {
'deeplabv3': (DeepLabHead, DeepLabV3),
'fcn': (FCNHead, FCN),
}
layer4 = list(backbone.layer4)[-1]
inplanes = list(layer4.children())[-2].num_features
classifier = model_map[name][0](inplanes, num_classes)
base_model = model_map[name][1]
model = ColourTransferModel(
base_model(backbone, classifier, aux_classifier), **kwargs)
return model
def __init__(self, cfg=None, phase='train'):
super(RetinaFace, self).__init__()
self.phase = phase
backbone = None
if cfg.MODEL.backbone == "mobilenet0.25":
backbone = MobileNetV1()
if cfg.TRAIN.pretrained:
ckpt = torch.load(
'./pretrained_weights/mobilenetv1x0.25_pretrain.pth',
map_location=torch.device('cpu'))
state_dict = OrderedDict()
for k, v in ckpt['state_dict'].items():
name = k[7:] # remove module
state_dict[name] = v
backbone.load_state_dict(state_dict)
elif cfg.MODEL.backbone == "resnet50":
backbone = torchvision.models.resnet50(
pretrained=cfg['pretrained'])
self.features = IntermediateLayerGetter(backbone,
cfg.MODEL.return_layers)
in_channels_list = cfg.MODEL.in_channels_list
out_channels = cfg.MODEL.out_channels
self.fpn = FPN(in_channels_list, out_channels)
self.ssh1 = SSH(out_channels, out_channels)
self.ssh2 = SSH(out_channels, out_channels)
self.ssh3 = SSH(out_channels, out_channels)
self.cls_heads = self._make_cls_head(3, in_channels=out_channels)
self.bbox_heads = self._make_bbox_head(3, in_channels=out_channels)
self.ldmk_heads = self._make_ldmk_head(3, in_channels=out_channels)
self._init_weights()
def __init__(self,
backbone: nn.Module,
train_backbone: bool,
num_channels: int,
return_interm_layers: bool,
net_type="resnet"):
super().__init__()
for name, parameter in backbone.named_parameters():
if not train_backbone or 'layer2' not in name and 'layer3' not in name:
parameter.requires_grad_(
False
) # here should allow users to specify which layers to freeze !
# print('freeze %s'%name)
if return_interm_layers:
if net_type == "resnet":
return_layers = {
"layer1": "0",
"layer2": "1",
"layer3": "2"
} # stride = 4, 8, 16
elif net_type == "repvgg":
return_layers = {"stage1": "0", "stage2": "1", "stage3": "2"}
else:
raise ValueError()
else:
if net_type == "resnet":
return_layers = {'layer3': "0"} # stride = 16
elif net_type == "repvgg":
return_layers = {'stage3': "0"} # stride = 16
else:
raise ValueError()
self.body = IntermediateLayerGetter(
backbone, return_layers=return_layers) # method in torchvision
self.num_channels = num_channels
def __init__(self,
arch,
pretrained_backbone=True,
output_stride=16,
num_classes=20):
super(DeepLabv3PlusModel, self).__init__()
if output_stride == 16:
replace_stride_with_dilation = [False, False, True]
atrous_rates = [6, 12, 18]
elif output_stride == 8:
replace_stride_with_dilation = [False, True, True]
atrous_rates = [12, 24, 36]
else:
raise ValueError('output_stride can be 8 or 16.')
backbone = create_resnet_backbone(
arch,
pretrained=pretrained_backbone,
replace_stride_with_dilation=replace_stride_with_dilation)
return_layers = {'layer1': 'low', 'layer4': 'high'}
self.backbone = IntermediateLayerGetter(backbone,
return_layers=return_layers)
self.classifier = DeepLabv3PlusHead(RESNET_STRIDE4_FEATURE_SIZES[arch],
48,
RESNET_HIGH_FEATURE_SIZES[arch],
num_classes, atrous_rates)
def _sar(arch: str,
pretrained: bool,
backbone_fn: Callable[[bool], nn.Module],
layer: str,
pretrained_backbone: bool = True,
**kwargs: Any) -> SAR:
pretrained_backbone = pretrained_backbone and not pretrained
# Patch the config
_cfg = deepcopy(default_cfgs[arch])
_cfg['vocab'] = kwargs.get('vocab', _cfg['vocab'])
_cfg['input_shape'] = kwargs.get('input_shape', _cfg['input_shape'])
# Feature extractor
feat_extractor = IntermediateLayerGetter(
backbone_fn(pretrained_backbone),
{layer: 'features'},
)
kwargs['vocab'] = _cfg['vocab']
kwargs['input_shape'] = _cfg['input_shape']
# Build the model
model = SAR(feat_extractor, cfg=_cfg, **kwargs)
# Load pretrained parameters
if pretrained:
load_pretrained_params(model, default_cfgs[arch]['url'])
return model
def _segm_resnet_emb(name,
backbone_name,
num_classes,
aux,
args,
pretrained_backbone=True,
embedding_model='mean',
aspp_option='v3'):
backbone = resnet.__dict__[backbone_name](
pretrained=pretrained_backbone,
replace_stride_with_dilation=[False, True, True])
return_layers = {'layer4': 'out'}
if aux:
return_layers['layer3'] = 'aux'
backbone = IntermediateLayerGetter(backbone, return_layers=return_layers)
aux_classifier = None
if aux:
inplanes = 1024
aux_classifier = FCNHead(inplanes, num_classes)
model_map = {'deeplabv3': (DeepLabHeadEmb, DeepLabV3Emb)}
inplanes = 2048
classifier = model_map[name][0](inplanes, num_classes, aspp_option, args)
base_model = model_map[name][1]
model = base_model(backbone, classifier, aux_classifier)
return model
