| 207 | |
| 208 | |
| 209 | class SLIP(CLIP): |
| 210 | def __init__(self,image_encode, text_encode, use_allgather, |
| 211 | EDA=True, feature_dim=1024, sim_dim=256, forward_type='split', return_sim=False): |
| 212 | super(SLIP, self).__init__(image_encode, text_encode, use_allgather) |
| 213 | self.return_sim = return_sim |
| 214 | if self.return_sim: |
| 215 | self.predictor_sim = projection_MLP(feature_dim, hidden_dim=4096, out_dim=sim_dim, out_bn=False) |
| 216 | self.forward_type = forward_type |
| 217 | |
| 218 | def text_modules(self): |
| 219 | ret = super(SLIP, self).text_modules() |
| 220 | return ret |
| 221 | |
| 222 | def visual_modules(self): |
| 223 | ret = super(SLIP, self).visual_modules() |
| 224 | ret.extend([self.predictor_sim]) |
| 225 | return ret |
| 226 | |
| 227 | def encode_image(self, image, return_dense=False, return_sim=False): |
| 228 | if return_dense: |
| 229 | output = self.visual(image.type(self.dtype), return_dense=return_dense, return_feature=return_sim) |
| 230 | else: |
| 231 | output = self.visual(image.type(self.dtype), return_feature=return_sim) |
| 232 | if return_sim: |
| 233 | sim_out = self.predictor_sim(output[-1]) |
| 234 | output = (*output[:-1], sim_out) # change image feature |
| 235 | return output |
| 236 | |
| 237 | def encode_text(self, text, text_mask_type=None, return_sim=False): |
| 238 | assert not return_sim |
| 239 | if text_mask_type: |
| 240 | output = self.text_encoder(text, mask_type=text_mask_type) |
| 241 | else: |
| 242 | output = self.text_encoder(text) |
| 243 | return output |
| 244 | |
| 245 | def forward(self, input, return_dict=False): |
| 246 | # input |
| 247 | images = input['images'] |
| 248 | images_base, images_1, images_2 = torch.split(images, [3,3,3], dim=1) |
| 249 | texts = input['captions'] |
| 250 | texts = self.sample_captions(texts) |
| 251 | text_features = self.encode_text(texts) |
| 252 | |
| 253 | image_features = self.encode_image(images_base) |
| 254 | image_features_1, image_sim_1 = self.encode_image(images_1, return_sim=True) |
| 255 | image_features_2, image_sim_2 = self.encode_image(images_2, return_sim=True) |
| 256 | |
| 257 | # normalized features |
| 258 | image_features = image_features / (image_features.norm(dim=-1, keepdim=True)) |
| 259 | text_features = text_features / (text_features.norm(dim=-1, keepdim=True)+1e-10) |
| 260 | |
| 261 | # image_sim_1 = image_sim_1 / (image_sim_1.norm(dim=-1, keepdim=True)) |
| 262 | # image_sim_2 = image_sim_2 / (image_sim_2.norm(dim=-1, keepdim=True)) |
| 263 | |
| 264 | # cosine similarity as logits |
| 265 | logit_scale = self.logit_scale.exp() |
| 266 | if self.training and self.use_allgather: |
no outgoing calls
no test coverage detected