Fixing hieradet (sam2) tests

This commit is contained in:
Ross Wightman 2024-08-16 14:33:40 -07:00
parent e035381171
commit 0b05122cda
2 changed files with 6 additions and 5 deletions

View File

@ -52,14 +52,14 @@ FEAT_INTER_FILTERS = [
'vision_transformer', 'vision_transformer_sam', 'vision_transformer_hybrid', 'vision_transformer_relpos',
'beit', 'mvitv2', 'eva', 'cait', 'xcit', 'volo', 'twins', 'deit', 'swin_transformer', 'swin_transformer_v2',
'swin_transformer_v2_cr', 'maxxvit', 'efficientnet', 'mobilenetv3', 'levit', 'efficientformer', 'resnet',
'regnet', 'byobnet', 'byoanet', 'mlp_mixer', 'hiera', 'fastvit',
'regnet', 'byobnet', 'byoanet', 'mlp_mixer', 'hiera', 'fastvit', 'hieradet_sam2'
]
# transformer / hybrid models don't support full set of spatial / feature APIs and/or have spatial output.
NON_STD_FILTERS = [
'vit_*', 'tnt_*', 'pit_*', 'coat_*', 'cait_*', '*mixer_*', 'gmlp_*', 'resmlp_*', 'twins_*',
'convit_*', 'levit*', 'visformer*', 'deit*', 'xcit_*', 'crossvit_*', 'beit*',
'poolformer_*', 'volo_*', 'sequencer2d_*', 'mvitv2*', 'gcvit*', 'efficientformer*',
'poolformer_*', 'volo_*', 'sequencer2d_*', 'mvitv2*', 'gcvit*', 'efficientformer*', 'sam_hiera*',
'eva_*', 'flexivit*', 'eva02*', 'samvit_*', 'efficientvit_m*', 'tiny_vit_*', 'hiera_*', 'vitamin*', 'test_vit*',
]
NUM_NON_STD = len(NON_STD_FILTERS)

View File

@ -45,7 +45,7 @@ def window_unpartition(windows: torch.Tensor, window_size: Tuple[int, int], hw:
"""
H, W = hw
B = windows.shape[0] // (H * W // window_size[0] // window_size[1])
x = windows.view(B, H // window_size[0], W // window_size[0], window_size[0], window_size[1], -1)
x = windows.view(B, H // window_size[0], W // window_size[1], window_size[0], window_size[1], -1)
x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
return x
@ -567,11 +567,12 @@ def checkpoint_filter_fn(state_dict, model=None, prefix=''):
def _create_hiera_det(variant: str, pretrained: bool = False, **kwargs) -> HieraDet:
out_indices = kwargs.pop('out_indices', 4)
if True: # kwargs.get('pretrained_cfg', '') == '?':
checkpoint_prefix = ''
if 'sam2' in variant:
# SAM2 pretrained weights have no classifier or final norm-layer (`head.norm`)
# This is workaround loading with num_classes=0 w/o removing norm-layer.
kwargs.setdefault('pretrained_strict', False)
checkpoint_prefix = 'image_encoder.trunk.' if 'sam2' in variant else ''
checkpoint_prefix = 'image_encoder.trunk.'
return build_model_with_cfg(
HieraDet,
variant,