diff --git a/timm/models/resnetv2.py b/timm/models/resnetv2.py index 6611ae49..3ce0605a 100644 --- a/timm/models/resnetv2.py +++ b/timm/models/resnetv2.py @@ -8,7 +8,9 @@ Additionally, supports non pre-activation bottleneck for use as a backbone for V extra padding support to allow porting of official Hybrid ResNet pretrained weights from https://github.com/google-research/vision_transformer -Thanks to the Google team for the above two repositories and associated papers. +Thanks to the Google team for the above two repositories and associated papers: +* Big Transfer (BiT): General Visual Representation Learning - https://arxiv.org/abs/1912.11370 +* An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale - https://arxiv.org/abs/2010.11929 Original copyright of Google code below, modifications by Ross Wightman, Copyright 2020. """ @@ -86,19 +88,19 @@ default_cfgs = { num_classes=21843), - # trained on imagenet-1k - 'resnetv2_50x1_bits': _cfg( - url='https://storage.googleapis.com/bit_models/BiT-S-R50x1-ILSVRC2012.npz'), - 'resnetv2_50x3_bits': _cfg( - url='https://storage.googleapis.com/bit_models/BiT-S-R50x3-ILSVRC2012.npz'), - 'resnetv2_101x1_bits': _cfg( - url='https://storage.googleapis.com/bit_models/BiT-S-R101x3-ILSVRC2012.npz'), - 'resnetv2_101x3_bits': _cfg( - url='https://storage.googleapis.com/bit_models/BiT-S-R101x3-ILSVRC2012.npz'), - 'resnetv2_152x2_bits': _cfg( - url='https://storage.googleapis.com/bit_models/BiT-S-R152x2-ILSVRC2012.npz'), - 'resnetv2_152x4_bits': _cfg( - url='https://storage.googleapis.com/bit_models/BiT-S-R152x4-ILSVRC2012.npz'), + # trained on imagenet-1k, NOTE not overly interesting set of weights, leaving disabled for now + # 'resnetv2_50x1_bits': _cfg( + # url='https://storage.googleapis.com/bit_models/BiT-S-R50x1.npz'), + # 'resnetv2_50x3_bits': _cfg( + # url='https://storage.googleapis.com/bit_models/BiT-S-R50x3.npz'), + # 'resnetv2_101x1_bits': _cfg( + # url='https://storage.googleapis.com/bit_models/BiT-S-R101x3.npz'), + # 'resnetv2_101x3_bits': _cfg( + # url='https://storage.googleapis.com/bit_models/BiT-S-R101x3.npz'), + # 'resnetv2_152x2_bits': _cfg( + # url='https://storage.googleapis.com/bit_models/BiT-S-R152x2.npz'), + # 'resnetv2_152x4_bits': _cfg( + # url='https://storage.googleapis.com/bit_models/BiT-S-R152x4.npz'), } @@ -358,8 +360,8 @@ class ResNetV2(nn.Module): self.feature_info = [] stem_chs = make_div(stem_chs * wf) self.stem = create_stem(in_chans, stem_chs, stem_type, preact, conv_layer=conv_layer, norm_layer=norm_layer) - if not preact: - self.feature_info.append(dict(num_chs=stem_chs, reduction=4, module='stem')) + # NOTE no, reduction 2 feature if preact + self.feature_info.append(dict(num_chs=stem_chs, reduction=2, module='' if preact else 'stem.norm')) prev_chs = stem_chs curr_stride = 4 @@ -372,21 +374,19 @@ class ResNetV2(nn.Module): if curr_stride >= output_stride: dilation *= stride stride = 1 - if preact: - self.feature_info += [dict(num_chs=prev_chs, reduction=curr_stride, module=f'stages.{stage_idx}.norm1')] stage = ResNetStage( prev_chs, out_chs, stride=stride, dilation=dilation, depth=d, avg_down=avg_down, act_layer=act_layer, conv_layer=conv_layer, norm_layer=norm_layer, block_fn=block_fn) prev_chs = out_chs curr_stride *= stride - if not preact: - self.feature_info += [dict(num_chs=prev_chs, reduction=curr_stride, module=f'stages.{stage_idx}')] + feat_name = f'stages.{stage_idx}' + if preact: + feat_name = f'stages.{stage_idx + 1}.blocks.0.norm1' if (stage_idx + 1) != len(channels) else 'norm' + self.feature_info += [dict(num_chs=prev_chs, reduction=curr_stride, module=feat_name)] self.stages.add_module(str(stage_idx), stage) self.num_features = prev_chs self.norm = norm_layer(self.num_features) if preact else nn.Identity() - if preact: - self.feature_info += [dict(num_chs=self.num_features, reduction=curr_stride, module=f'norm')] self.head = ClassifierHead( self.num_features, num_classes, pool_type=global_pool, drop_rate=self.drop_rate, use_conv=True) @@ -446,9 +446,15 @@ class ResNetV2(nn.Module): def _create_resnetv2(variant, pretrained=False, **kwargs): # FIXME feature map extraction is not setup properly for pre-activation mode right now + preact = kwargs.get('preact', True) + feature_cfg = dict(flatten_sequential=True) + if preact: + feature_cfg['feature_cls'] = 'hook' + feature_cfg['out_indices'] = (1, 2, 3, 4) # no stride 2, 0 level feat for preact + return build_model_with_cfg( ResNetV2, variant, pretrained, default_cfg=default_cfgs[variant], pretrained_custom_load=True, - feature_cfg=dict(flatten_sequential=True), **kwargs) + feature_cfg=feature_cfg, **kwargs) @register_model @@ -496,83 +502,85 @@ def resnetv2_152x4_bitm(pretrained=False, **kwargs): @register_model def resnetv2_50x1_bitm_in21k(pretrained=False, **kwargs): return _create_resnetv2( - 'resnetv2_50x1_bitm', pretrained=pretrained, + 'resnetv2_50x1_bitm_in21k', pretrained=pretrained, num_classes=kwargs.get('num_classes', 21843), layers=[3, 4, 6, 3], width_factor=1, stem_type='fixed', **kwargs) @register_model def resnetv2_50x3_bitm_in21k(pretrained=False, **kwargs): return _create_resnetv2( - 'resnetv2_50x3_bitm', pretrained=pretrained, + 'resnetv2_50x3_bitm_in21k', pretrained=pretrained, num_classes=kwargs.get('num_classes', 21843), layers=[3, 4, 6, 3], width_factor=3, stem_type='fixed', **kwargs) @register_model def resnetv2_101x1_bitm_in21k(pretrained=False, **kwargs): return _create_resnetv2( - 'resnetv2_101x1_bitm', pretrained=pretrained, + 'resnetv2_101x1_bitm_in21k', pretrained=pretrained, num_classes=kwargs.get('num_classes', 21843), layers=[3, 4, 23, 3], width_factor=1, stem_type='fixed', **kwargs) @register_model def resnetv2_101x3_bitm_in21k(pretrained=False, **kwargs): return _create_resnetv2( - 'resnetv2_101x3_bitm', pretrained=pretrained, + 'resnetv2_101x3_bitm_in21k', pretrained=pretrained, num_classes=kwargs.get('num_classes', 21843), layers=[3, 4, 23, 3], width_factor=3, stem_type='fixed', **kwargs) @register_model def resnetv2_152x2_bitm_in21k(pretrained=False, **kwargs): return _create_resnetv2( - 'resnetv2_152x2_bitm', pretrained=pretrained, + 'resnetv2_152x2_bitm_in21k', pretrained=pretrained, num_classes=kwargs.get('num_classes', 21843), layers=[3, 8, 36, 3], width_factor=2, stem_type='fixed', **kwargs) @register_model def resnetv2_152x4_bitm_in21k(pretrained=False, **kwargs): return _create_resnetv2( - 'resnetv2_152x4_bitm', pretrained=pretrained, + 'resnetv2_152x4_bitm_in21k', pretrained=pretrained, num_classes=kwargs.get('num_classes', 21843), layers=[3, 8, 36, 3], width_factor=4, stem_type='fixed', **kwargs) -@register_model -def resnetv2_50x1_bits(pretrained=False, **kwargs): - return _create_resnetv2( - 'resnetv2_50x1_bits', pretrained=pretrained, - layers=[3, 4, 6, 3], width_factor=1, stem_type='fixed', **kwargs) - - -@register_model -def resnetv2_50x3_bits(pretrained=False, **kwargs): - return _create_resnetv2( - 'resnetv2_50x3_bits', pretrained=pretrained, - layers=[3, 4, 6, 3], width_factor=3, stem_type='fixed', **kwargs) - - -@register_model -def resnetv2_101x1_bits(pretrained=False, **kwargs): - return _create_resnetv2( - 'resnetv2_101x1_bits', pretrained=pretrained, - layers=[3, 4, 23, 3], width_factor=1, stem_type='fixed', **kwargs) - - -@register_model -def resnetv2_101x3_bits(pretrained=False, **kwargs): - return _create_resnetv2( - 'resnetv2_101x3_bits', pretrained=pretrained, - layers=[3, 4, 23, 3], width_factor=3, stem_type='fixed', **kwargs) - - -@register_model -def resnetv2_152x2_bits(pretrained=False, **kwargs): - return _create_resnetv2( - 'resnetv2_152x2_bits', pretrained=pretrained, - layers=[3, 8, 36, 3], width_factor=2, stem_type='fixed', **kwargs) - - -@register_model -def resnetv2_152x4_bits(pretrained=False, **kwargs): - return _create_resnetv2( - 'resnetv2_152x4_bits', pretrained=pretrained, - layers=[3, 8, 36, 3], width_factor=4, stem_type='fixed', **kwargs) +# NOTE the 'S' versions of the model weights arent as interesting as original 21k or transfer to 1K M. +# @register_model +# def resnetv2_50x1_bits(pretrained=False, **kwargs): +# return _create_resnetv2( +# 'resnetv2_50x1_bits', pretrained=pretrained, +# layers=[3, 4, 6, 3], width_factor=1, stem_type='fixed', **kwargs) +# +# +# @register_model +# def resnetv2_50x3_bits(pretrained=False, **kwargs): +# return _create_resnetv2( +# 'resnetv2_50x3_bits', pretrained=pretrained, +# layers=[3, 4, 6, 3], width_factor=3, stem_type='fixed', **kwargs) +# +# +# @register_model +# def resnetv2_101x1_bits(pretrained=False, **kwargs): +# return _create_resnetv2( +# 'resnetv2_101x1_bits', pretrained=pretrained, +# layers=[3, 4, 23, 3], width_factor=1, stem_type='fixed', **kwargs) +# +# +# @register_model +# def resnetv2_101x3_bits(pretrained=False, **kwargs): +# return _create_resnetv2( +# 'resnetv2_101x3_bits', pretrained=pretrained, +# layers=[3, 4, 23, 3], width_factor=3, stem_type='fixed', **kwargs) +# +# +# @register_model +# def resnetv2_152x2_bits(pretrained=False, **kwargs): +# return _create_resnetv2( +# 'resnetv2_152x2_bits', pretrained=pretrained, +# layers=[3, 8, 36, 3], width_factor=2, stem_type='fixed', **kwargs) +# +# +# @register_model +# def resnetv2_152x4_bits(pretrained=False, **kwargs): +# return _create_resnetv2( +# 'resnetv2_152x4_bits', pretrained=pretrained, +# layers=[3, 8, 36, 3], width_factor=4, stem_type='fixed', **kwargs) +# diff --git a/timm/models/vision_transformer.py b/timm/models/vision_transformer.py index 9b96e04e..ff5bd676 100644 --- a/timm/models/vision_transformer.py +++ b/timm/models/vision_transformer.py @@ -79,23 +79,27 @@ default_cfgs = { # patch models, imagenet21k (weights ported from official JAX impl) 'vit_base_patch16_224_in21k': _cfg( - url='', + url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_patch16_224_in21k-e5005f0a.pth', num_classes=21843, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)), 'vit_base_patch32_224_in21k': _cfg( - url='', + url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_patch32_224_in21k-8db57226.pth', num_classes=21843, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)), 'vit_large_patch16_224_in21k': _cfg( - url='', + url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_patch16_224_in21k-606da67d.pth', num_classes=21843, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)), 'vit_large_patch32_224_in21k': _cfg( - url='', + url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_patch32_224_in21k-9046d2e7.pth', num_classes=21843, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)), 'vit_huge_patch14_224_in21k': _cfg( - url='', + url='', # FIXME I have weights for this but > 2GB limit for github release binaries num_classes=21843, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)), # hybrid models (weights ported from official JAX impl) + 'vit_base_resnet50_224_in21k': _cfg( + url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_resnet50_224_in21k-6f7c7740.pth', + mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=0.9), 'vit_base_resnet50_384': _cfg( + url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_resnet50_384-9fd3c705.pth', input_size=(3, 384, 384), mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=1.0), # hybrid models (my experiments) @@ -269,6 +273,7 @@ class VisionTransformer(nn.Module): # Representation layer if representation_size: + self.num_features = representation_size self.pre_logits = nn.Sequential(OrderedDict([ ('fc', nn.Linear(embed_dim, representation_size)), ('act', nn.Tanh()) @@ -315,12 +320,12 @@ class VisionTransformer(nn.Module): for blk in self.blocks: x = blk(x) - x = self.norm(x) - return x[:, 0] + x = self.norm(x)[:, 0] + x = self.pre_logits(x) + return x def forward(self, x): x = self.forward_features(x) - x = self.pre_logits(x) x = self.head(x) return x @@ -407,7 +412,7 @@ def vit_large_patch16_224(pretrained=False, **kwargs): @register_model def vit_large_patch32_224(pretrained=False, **kwargs): model = VisionTransformer( - img_size=224, patch_size=32, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True, + img_size=224, patch_size=32, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs) model.default_cfg = default_cfgs['vit_large_patch32_224'] if pretrained: @@ -418,7 +423,7 @@ def vit_large_patch32_224(pretrained=False, **kwargs): @register_model def vit_large_patch16_384(pretrained=False, **kwargs): model = VisionTransformer( - img_size=384, patch_size=16, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True, + img_size=384, patch_size=16, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs) model.default_cfg = default_cfgs['vit_large_patch16_384'] if pretrained: @@ -426,22 +431,12 @@ def vit_large_patch16_384(pretrained=False, **kwargs): return model -@register_model -def vit_large_patch32_384(pretrained=False, **kwargs): - model = VisionTransformer( - img_size=384, patch_size=32, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True, - norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs) - model.default_cfg = default_cfgs['vit_large_patch32_384'] - if pretrained: - load_pretrained(model, num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3)) - return model - - @register_model def vit_base_patch16_224_in21k(pretrained=False, **kwargs): + num_classes = kwargs.get('num_classes', 21843) model = VisionTransformer( - patch_size=16, num_classes=21843, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True, - norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs) + patch_size=16, num_classes=num_classes, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True, + representation_size=768, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs) model.default_cfg = default_cfgs['vit_base_patch16_224_in21k'] if pretrained: load_pretrained( @@ -451,9 +446,10 @@ def vit_base_patch16_224_in21k(pretrained=False, **kwargs): @register_model def vit_base_patch32_224_in21k(pretrained=False, **kwargs): + num_classes = kwargs.get('num_classes', 21843) model = VisionTransformer( - img_size=224, num_classes=21843, patch_size=32, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, - qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs) + img_size=224, num_classes=num_classes, patch_size=32, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, + qkv_bias=True, representation_size=768, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs) model.default_cfg = default_cfgs['vit_base_patch32_224_in21k'] if pretrained: load_pretrained(model, num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3)) @@ -462,9 +458,10 @@ def vit_base_patch32_224_in21k(pretrained=False, **kwargs): @register_model def vit_large_patch16_224_in21k(pretrained=False, **kwargs): + num_classes = kwargs.get('num_classes', 21843) model = VisionTransformer( - patch_size=16, num_classes=21843, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True, - norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs) + patch_size=16, num_classes=num_classes, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True, + representation_size=1024, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs) model.default_cfg = default_cfgs['vit_large_patch16_224_in21k'] if pretrained: load_pretrained(model, num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3)) @@ -473,9 +470,10 @@ def vit_large_patch16_224_in21k(pretrained=False, **kwargs): @register_model def vit_large_patch32_224_in21k(pretrained=False, **kwargs): + num_classes = kwargs.get('num_classes', 21843) model = VisionTransformer( - img_size=224, num_classes=21843, patch_size=32, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True, - norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs) + img_size=224, num_classes=num_classes, patch_size=32, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, + qkv_bias=True, representation_size=1024, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs) model.default_cfg = default_cfgs['vit_large_patch32_224_in21k'] if pretrained: load_pretrained(model, num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3)) @@ -484,15 +482,31 @@ def vit_large_patch32_224_in21k(pretrained=False, **kwargs): @register_model def vit_huge_patch14_224_in21k(pretrained=False, **kwargs): + num_classes = kwargs.get('num_classes', 21843) model = VisionTransformer( - img_size=224, patch_size=14, num_classes=21843, embed_dim=1280, depth=32, num_heads=16, mlp_ratio=4, - qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs) + img_size=224, patch_size=14, num_classes=num_classes, embed_dim=1280, depth=32, num_heads=16, mlp_ratio=4, + qkv_bias=True, representation_size=1280, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs) model.default_cfg = default_cfgs['vit_huge_patch14_224_in21k'] if pretrained: load_pretrained(model, num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3)) return model +@register_model +def vit_base_resnet50_224_in21k(pretrained=False, **kwargs): + # create a ResNetV2 w/o pre-activation, that uses StdConv and GroupNorm and has 3 stages, no head + num_classes = kwargs.get('num_classes', 21843) + backbone = ResNetV2( + layers=(3, 4, 9), preact=False, stem_type='same', conv_layer=StdConv2dSame, num_classes=0, global_pool='') + model = VisionTransformer( + img_size=224, num_classes=num_classes, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True, + hybrid_backbone=backbone, representation_size=768, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs) + model.default_cfg = default_cfgs['vit_base_resnet50_224_in21k'] + if pretrained: + load_pretrained(model, num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3)) + return model + + @register_model def vit_base_resnet50_384(pretrained=False, **kwargs): # create a ResNetV2 w/o pre-activation, that uses StdConv and GroupNorm and has 3 stages, no head diff --git a/validate.py b/validate.py index 645dfd1d..4eedd6fb 100755 --- a/validate.py +++ b/validate.py @@ -60,7 +60,7 @@ parser.add_argument('--std', type=float, nargs='+', default=None, metavar='STD' help='Override std deviation of of dataset') parser.add_argument('--interpolation', default='', type=str, metavar='NAME', help='Image resize interpolation type (overrides model)') -parser.add_argument('--num-classes', type=int, default=1000, +parser.add_argument('--num-classes', type=int, default=None, help='Number classes in dataset') parser.add_argument('--class-map', default='', type=str, metavar='FILENAME', help='path to class to idx mapping file (default: "")')