(self, pretrained_dict=None, pretrained_layers=[], verbose=True)
| 476 | raise TypeError('pretrained must be a str or None') |
| 477 | |
| 478 | def load_weights(self, pretrained_dict=None, pretrained_layers=[], verbose=True): |
| 479 | model_dict = self.state_dict() |
| 480 | |
| 481 | missed_dict = [k for k in model_dict.keys() if k not in pretrained_dict] |
| 482 | logger.info(f'=> Missed keys {missed_dict}') |
| 483 | unexpected_dict = [k for k in pretrained_dict.keys() if k not in model_dict] |
| 484 | logger.info(f'=> Unexpected keys {unexpected_dict}') |
| 485 | |
| 486 | pretrained_dict = { |
| 487 | k: v for k, v in pretrained_dict.items() |
| 488 | if k in model_dict.keys() |
| 489 | } |
| 490 | |
| 491 | need_init_state_dict = {} |
| 492 | for k, v in pretrained_dict.items(): |
| 493 | need_init = ( |
| 494 | ( |
| 495 | k.split('.')[0] in pretrained_layers |
| 496 | or pretrained_layers[0] == '*' |
| 497 | ) |
| 498 | and 'relative_position_index' not in k |
| 499 | and 'attn_mask' not in k |
| 500 | ) |
| 501 | |
| 502 | if need_init: |
| 503 | # if verbose: |
| 504 | # logger.info(f'=> init {k} from {pretrained}') |
| 505 | |
| 506 | if ('pool_layers' in k) or ('focal_layers' in k) and v.size() != model_dict[k].size(): |
| 507 | table_pretrained = v |
| 508 | table_current = model_dict[k] |
| 509 | fsize1 = table_pretrained.shape[2] |
| 510 | fsize2 = table_current.shape[2] |
| 511 | |
| 512 | # NOTE: different from interpolation used in self-attention, we use padding or clipping for focal conv |
| 513 | if fsize1 < fsize2: |
| 514 | table_pretrained_resized = torch.zeros(table_current.shape) |
| 515 | table_pretrained_resized[:, :, (fsize2-fsize1)//2:-(fsize2-fsize1)//2, (fsize2-fsize1)//2:-(fsize2-fsize1)//2] = table_pretrained |
| 516 | v = table_pretrained_resized |
| 517 | elif fsize1 > fsize2: |
| 518 | table_pretrained_resized = table_pretrained[:, :, (fsize1-fsize2)//2:-(fsize1-fsize2)//2, (fsize1-fsize2)//2:-(fsize1-fsize2)//2] |
| 519 | v = table_pretrained_resized |
| 520 | |
| 521 | |
| 522 | if ("modulation.f" in k or "pre_conv" in k): |
| 523 | table_pretrained = v |
| 524 | table_current = model_dict[k] |
| 525 | if table_pretrained.shape != table_current.shape: |
| 526 | if len(table_pretrained.shape) == 2: |
| 527 | dim = table_pretrained.shape[1] |
| 528 | assert table_current.shape[1] == dim |
| 529 | L1 = table_pretrained.shape[0] |
| 530 | L2 = table_current.shape[0] |
| 531 | |
| 532 | if L1 < L2: |
| 533 | table_pretrained_resized = torch.zeros(table_current.shape) |
| 534 | # copy for linear project |
| 535 | table_pretrained_resized[:2*dim] = table_pretrained[:2*dim] |
no test coverage detected