| 252 | |
| 253 | @torch.no_grad() |
| 254 | def load_pretrained(self, state_dict, lookup_22k): |
| 255 | |
| 256 | new_state_dict = {} |
| 257 | for state_key, state_value in state_dict.items(): |
| 258 | keys = state_key.split('.') |
| 259 | m = self |
| 260 | for key in keys: |
| 261 | if key.isdigit(): |
| 262 | m = m[int(key)] |
| 263 | else: |
| 264 | m = getattr(m, key) |
| 265 | if m.shape == state_value.shape: |
| 266 | new_state_dict[state_key] = state_value |
| 267 | else: |
| 268 | # Ignore different shapes |
| 269 | if 'relative_position_index' in keys: |
| 270 | new_state_dict[state_key] = m.data |
| 271 | if 'q_grid' in keys: |
| 272 | new_state_dict[state_key] = m.data |
| 273 | if 'reference' in keys: |
| 274 | new_state_dict[state_key] = m.data |
| 275 | # Bicubic Interpolation |
| 276 | if 'relative_position_bias_table' in keys: |
| 277 | n, c = state_value.size() |
| 278 | l_side = int(math.sqrt(n)) |
| 279 | assert n == l_side ** 2 |
| 280 | L = int(math.sqrt(m.shape[0])) |
| 281 | pre_interp = state_value.reshape(1, l_side, l_side, c).permute(0, 3, 1, 2) |
| 282 | post_interp = F.interpolate(pre_interp, (L, L), mode='bicubic') |
| 283 | new_state_dict[state_key] = post_interp.reshape(c, L ** 2).permute(1, 0) |
| 284 | if 'rpe_table' in keys: |
| 285 | c, h, w = state_value.size() |
| 286 | C, H, W = m.data.size() |
| 287 | pre_interp = state_value.unsqueeze(0) |
| 288 | post_interp = F.interpolate(pre_interp, (H, W), mode='bicubic') |
| 289 | new_state_dict[state_key] = post_interp.squeeze(0) |
| 290 | if 'cls_head' in keys: |
| 291 | new_state_dict[state_key] = state_value[lookup_22k] |
| 292 | |
| 293 | msg = self.load_state_dict(new_state_dict, strict=False) |
| 294 | return msg |
| 295 | |
| 296 | @torch.jit.ignore |
| 297 | def no_weight_decay(self): |