def init_weights(self):
org_resnet = model_zoo.load_url(model_urls[self.name])
# drop orginal resnet fc layer, add 'None' in case of no fc layer, that will raise error
org_resnet.pop('fc.weight', None)
org_resnet.pop('fc.bias', None)
self.load_state_dict(org_resnet)
print("Initialize resnet from model zoo")
def init_pose_net(pose_net, name):
org_resnet = model_zoo.load_url(model_urls[name])
# drop orginal resnet fc layer, add 'None' in case of no fc layer, that will raise error
org_resnet.pop('fc.weight', None)
org_resnet.pop('fc.bias', None)
pose_net.backbone.load_state_dict(org_resnet)
print("Init Network from model zoo")
def __init__(self, num_classes):
super(ResNet183CProtein, self).__init__(BasicBlock, [2, 2, 2, 2])
self.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
self.avgpool = nn.AdaptiveAvgPool2d(1)
self.fc = nn.Sequential(OrderedDict([
('bn1', nn.BatchNorm1d(512)),
('drop1', nn.Dropout(p=0.5)),
('linear1', nn.Linear(512, num_classes))
]))
def __init__(self, num_classes):
super(ResNet18MaxProtein, self).__init__(BasicBlock, [2, 2, 2, 2])
self.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
self.conv1_y = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3))
self.ada_maxpool = nn.AdaptiveMaxPool2d(1)
self.fc = nn.Sequential(OrderedDict([
('bn1', nn.BatchNorm1d(512)),
('drop1', nn.Dropout(p=0.5)),
('linear1', nn.Linear(512, num_classes))
]))
def resnet18(pretrained=False, **kwargs):
"""Constructs a ResNet-18 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(BasicBlock, [2, 2, 2, 2], net_type='resnet18', **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet18']),
strict=False)
return model
def resnet101(pretrained=False, **kwargs):
"""Constructs a ResNet-101 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(Bottleneck, [3, 4, 23, 3], net_type='resnet101', **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet101']),
strict=False)
return model
def fpn_34(pretrained=True):
model = ResNet(BasicBlock, [3, 4, 6, 3], num_classes=17)
if pretrained:
state_dict = model.state_dict()
pretrained_dict = model_zoo.load_url(model_urls['resnet34'])
pretrained_dict = {
key: pretrained_dict[key]
for key in pretrained_dict.keys() if 'fc' not in key
}
state_dict.update(pretrained_dict)
model.load_state_dict(state_dict)
return model
def fpn_152(pretrained=True):
model = ResNet(Bottleneck, [3, 8, 36, 3], num_classes=17)
if pretrained:
state_dict = model.state_dict()
pretrained_dict = model_zoo.load_url(model_urls['resnet152'])
pretrained_dict = {
key: pretrained_dict[key]
for key in pretrained_dict.keys() if 'fc' not in key
}
state_dict.update(pretrained_dict)
model.load_state_dict(state_dict)
return model
def init_pose_net(pose_net, cfg):
if cfg.from_model_zoo:
_, _, _, name = resnet_spec[cfg.num_layers]
org_resnet = model_zoo.load_url(model_urls[name])
# drop orginal resnet fc layer, add 'None' in case of no fc layer, that will raise error
org_resnet.pop('fc.weight', None)
org_resnet.pop('fc.bias', None)
pose_net.backbone.load_state_dict(org_resnet)
# print('Loading pretrained model from {}'.format(os.path.join(cfg.pretrained, model_file)))
else:
if os.path.exists(cfg.pretrained):
model = torch.load(cfg.pretrained)
pose_net.load_state_dict(model['network'])
print("Init Network from pretrained", cfg.pretrained)
def fpn_50(pretrained=True):
model = ResNet(Bottleneck, [3, 4, 6, 3], num_classes=17)
if pretrained:
# load model dictionary
model_dict = model.state_dict()
# load pretrained model
pretrained_dict = model_zoo.load_url(model_urls['resnet50'])
# update model dictionary using pretrained model without classifier layer
model_dict.update({
key: pretrained_dict[key]
for key in pretrained_dict.keys() if 'fc' not in key
})
model.load_state_dict(model_dict)
return model
def __init__(self, num_classes):
super(ResNet34Protein3C, self).__init__(BasicBlock, [3, 4, 6, 3])
self.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
self.ada_avgpool = nn.AdaptiveAvgPool2d(1)
self.ada_maxpool = nn.AdaptiveMaxPool2d(1)
self.fc = nn.Sequential(OrderedDict([
('bn1', nn.BatchNorm1d(1024)),
('drop1', nn.Dropout(p=0.5)),
('linear1', nn.Linear(1024, 512)),
('relu1', nn.ReLU()),
('bn2', nn.BatchNorm1d(512)),
('drop2', nn.Dropout(p=0.5)),
('linear2', nn.Linear(512, num_classes))
]))
def resnet18_planet(pretrained=False):
model = resnet18(False, num_classes=17)
if pretrained:
# load model dictionary
model_dict = model.state_dict()
# load pretrained model
pretrained_dict = model_zoo.load_url(model_urls['resnet18'])
# update model dictionary using pretrained model without classifier layer
model_dict.update({
key: pretrained_dict[key]
for key in pretrained_dict.keys() if 'fc' not in key
})
model.load_state_dict(model_dict)
return model
def __init__(self, num_classes):
super(ResNet50Protein, self).__init__(Bottleneck, [3, 4, 6, 3])
self.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
self.conv1_y = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3))
self.ada_avgpool = nn.AdaptiveAvgPool2d(1)
self.ada_maxpool = nn.AdaptiveMaxPool2d(1)
self.fc = nn.Sequential(OrderedDict([
('bn1', nn.BatchNorm1d(4096)),
('drop1', nn.Dropout(p=0.5)),
('linear1', nn.Linear(4096, 1024)),
('relu1', nn.ReLU()),
('bn2', nn.BatchNorm1d(1024)),
('drop2', nn.Dropout(p=0.5)),
('linear2', nn.Linear(1024, num_classes))
]))
def __init__(self, block=BasicBlock, layers=[2, 2, 2], num_classes=1000):
self.inplanes = 64
self.out_planes = 128
super(LomoNet, self).__init__()
self.layer1 = self._make_layer(block, 64, layers[0], stride=2)
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
load_state_dict(self, model_zoo.load_url(model_urls['resnet18']))
def initialize_backbone_from_modelzoo(
backbone, # type: ResNetBackbone,
resnet_num_layers, # type: int
image_channels, # type: int
):
assert image_channels == 3 or image_channels == 4
_, _, _, name = resnet_spec[resnet_num_layers]
org_resnet = model_zoo.load_url(model_urls[name])
# Drop orginal resnet fc layer, add 'None' in case of no fc layer, that will raise error
org_resnet.pop('fc.weight', None)
org_resnet.pop('fc.bias', None)
# Load the backbone
if image_channels is 3:
backbone.load_state_dict(org_resnet)
elif image_channels is 4:
# Modify the first conv
conv1_weight_old = org_resnet['conv1.weight']
conv1_weight = torch.zeros((64, 4, 7, 7))
conv1_weight[:, 0:3, :, :] = conv1_weight_old
avg_weight = conv1_weight_old.mean(dim=1, keepdim=False)
conv1_weight[:, 3, :, :] = avg_weight
org_resnet['conv1.weight'] = conv1_weight
# Load it
backbone.load_state_dict(org_resnet)
def __init__(self, ngpu):
ResNet.__init__(self, BasicBlock, [3, 4, 6, 3])
self.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
self.ngpu = ngpu
self.output = nn.Sequential(nn.Linear(512, 1), nn.Sigmoid())
