-
Notifications
You must be signed in to change notification settings - Fork 15
/
Copy pathresnet_simclr.py
160 lines (120 loc) · 6.02 KB
/
resnet_simclr.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
import torch
import torch.nn as nn
import torchvision.models as models
from exceptions import InvalidBackboneError
class ResNetSimCLR(nn.Module):
def __init__(self, base_model, out_dim):
super(ResNetSimCLR, self).__init__()
self.resnet_dict = {"resnet18": models.resnet18(pretrained=False, num_classes=out_dim),
"resnet50": models.resnet50(pretrained=False, num_classes=out_dim)}
self.backbone = self._get_basemodel(base_model)
dim_mlp = self.backbone.fc.in_features
# add mlp projection head
self.backbone.fc = nn.Sequential(nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), self.backbone.fc)
def _get_basemodel(self, model_name):
try:
model = self.resnet_dict[model_name]
except KeyError:
raise InvalidBackboneError(
"Invalid backbone architecture. Check the config file and pass one of: resnet18 or resnet50")
else:
return model
def forward(self, x):
return self.backbone(x)
class NormalSimCLRDownstream(nn.Module):
def __init__(self, base_model, out_dim, checkpoint, num_classes, load_late=False):
super(NormalSimCLRDownstream, self).__init__()
self.resnet_dict = {"resnet18": models.resnet18(pretrained=False, num_classes=out_dim),
"resnet50": models.resnet50(pretrained=False, num_classes=out_dim)}
self.backbone = self._get_basemodel(base_model)
self.dim_mlp = self.backbone.fc.in_features
self.checkpoint = checkpoint
self.num_classes = num_classes
self.load_late = load_late
# add mlp projection head
self.backbone.fc = nn.Sequential(nn.Linear(self.dim_mlp, self.dim_mlp), nn.ReLU(), self.backbone.fc)
self.init_and_load_trained_state_dict()
def _get_basemodel(self, model_name):
try:
model = self.resnet_dict[model_name]
except KeyError:
raise InvalidBackboneError(
"Invalid backbone architecture. Check the config file and pass one of: resnet18 or resnet50")
else:
return model
def forward(self, x):
return self.backbone(x)
def init_and_load_trained_state_dict(self):
"""load the pre-trained backbone weights"""
if not self.load_late:
log_bb = self.load_state_dict(self.checkpoint["state_dict"], strict=True)
# replace the projection MLP with a classification layer
self.backbone.fc = torch.nn.Linear(self.dim_mlp, self.num_classes)
if self.load_late:
log_bb = self.load_state_dict(self.checkpoint["model_weights"], strict=True)
# freeze all layers but the last fc
for name, param in self.named_parameters():
if name not in ['backbone.fc.weight', 'backbone.fc.bias']:
param.requires_grad = False
class DoubleResNetSimCLR(nn.Module):
def __init__(self, base_model, out_dim):
super(DoubleResNetSimCLR, self).__init__()
self.resnet_dict = {"resnet18": models.resnet18,
"resnet50": models.resnet50,}
self.backbone1 = self.resnet_dict.get(base_model)(pretrained=False, num_classes=out_dim)
self.backbone2 = self.resnet_dict.get(base_model)(pretrained=False, num_classes=out_dim)
dim_mlp1 = self.backbone1.fc.in_features
dim_mlp2 = self.backbone2.fc.in_features
# add mlp projection head
self.backbone1.fc = nn.Sequential(nn.Linear(dim_mlp1, dim_mlp1), nn.ReLU(), self.backbone1.fc)
self.backbone2.fc = nn.Sequential(nn.Linear(dim_mlp2, dim_mlp2), nn.ReLU(), self.backbone2.fc)
def _get_basemodel(self, model_name):
try:
model = self.resnet_dict[model_name]
except KeyError:
raise InvalidBackboneError(
"Invalid backbone architecture. Check the config file and pass one of: resnet18 or resnet50")
else:
return model
def forward(self, x):
return {"s1" : self.backbone1(x["s1"]), "s2" : self.backbone2(x["s2"])}
class DoubleResNetSimCLRDownstream(nn.Module):
"""concatenate outputs from two backbones and add one linear layer"""
def __init__(self, base_model, out_dim):
super(DoubleResNetSimCLRDownstream, self).__init__()
self.resnet_dict = {"resnet18": models.resnet18,
"resnet50": models.resnet50,}
self.backbone1 = self.resnet_dict.get(base_model)(pretrained=False, num_classes=out_dim)
self.backbone2 = self.resnet_dict.get(base_model)(pretrained=False, num_classes=out_dim)
dim_mlp1 = self.backbone1.fc.in_features
dim_mlp2 = self.backbone2.fc.in_features
# add final linear layer
self.fc = nn.Linear(dim_mlp1 + dim_mlp2, out_dim, bias=True)
self.backbone1.fc = nn.Identity()
self.backbone2.fc = nn.Identity()
def _get_basemodel(self, model_name):
try:
model = self.resnet_dict[model_name]
except KeyError:
raise InvalidBackboneError(
"Invalid backbone architecture. Check the config file and pass one of: resnet18 or resnet50")
else:
return model
def forward(self, x):
x1 = self.backbone1(x["s1"])
x2 = self.backbone2(x["s2"])
z = torch.cat([x1, x2], dim=1)
z = self.fc(z)
return z
def load_trained_state_dict(self, weights):
"""load the pre-trained backbone weights"""
# remove the MLP projection heads
for k in list(weights.keys()):
if k.startswith(('backbone1.fc', 'backbone2.fc')):
del weights[k]
log = self.load_state_dict(weights, strict=False)
assert log.missing_keys == ['fc.weight', 'fc.bias']
# freeze all layers but the last fc
for name, param in self.named_parameters():
if name not in ['fc.weight', 'fc.bias']:
param.requires_grad = False