diff --git a/timm/models/davit.py b/timm/models/davit.py index 8e84f7ed..0259ec9f 100644 --- a/timm/models/davit.py +++ b/timm/models/davit.py @@ -7,10 +7,6 @@ attention in each block. The attention mechanisms used are linear in complexity. DaViT model defs and weights adapted from https://github.com/dingmyu/davit, original copyright below - - - - """ # Copyright (c) 2022 Mingyu Ding # All rights reserved. @@ -442,7 +438,7 @@ class DaViT(nn.Module): ) for layer_id, item in enumerate(stage_param) ]) - self.main_blocks.add_module(f'stage_{stage_id}', stage) + self.stages.add_module(f'stage_{stage_id}', stage) self.feature_info += [dict(num_chs=self.embed_dims[stage_id], reduction = 2, module=f'stage_{stage_id}')] @@ -536,7 +532,7 @@ def checkpoint_filter_fn(state_dict, model): out_dict = {} for k, v in state_dict.items(): - k = k.replace('main_blocks.', 'main_blocks.stage_') + k = k.replace('main_blocks.', 'stages.stage_') k = k.replace('head.', 'head.fc.') out_dict[k] = v return out_dict @@ -570,13 +566,13 @@ def _cfg(url='', **kwargs): # not sure how this should be set up default_cfgs = generate_default_cfgs({ - -'davit_tiny.msft_in1k': _cfg( - url="https://github.com/fffffgggg54/pytorch-image-models/releases/download/checkpoint/davit_tiny_ed28dd55.pth.tar"), -'davit_small.msft_in1k': _cfg( - url="https://github.com/fffffgggg54/pytorch-image-models/releases/download/checkpoint/davit_small_d1ecf281.pth.tar"), -'davit_base.msft_in1k': _cfg( - url="https://github.com/fffffgggg54/pytorch-image-models/releases/download/checkpoint/davit_base_67d9ac26.pth.tar"), + # official microsoft weights from https://github.com/dingmyu/davit + 'davit_tiny.msft_in1k': _cfg( + url="https://github.com/fffffgggg54/pytorch-image-models/releases/download/checkpoint/davit_tiny_ed28dd55.pth.tar"), + 'davit_small.msft_in1k': _cfg( + url="https://github.com/fffffgggg54/pytorch-image-models/releases/download/checkpoint/davit_small_d1ecf281.pth.tar"), + 'davit_base.msft_in1k': _cfg( + url="https://github.com/fffffgggg54/pytorch-image-models/releases/download/checkpoint/davit_base_67d9ac26.pth.tar"), })