Update davit.py

timm/models/davit.py

@@ -548,12 +548,34 @@ class DaViT(nn.Module):
             global_pool = self.head.global_pool.pool_type
         self.head = ClassifierHead(self.num_features, num_classes, pool_type=global_pool, drop_rate=self.drop_rate)
     
-
+    def forward_network(self, x : Tensor):
-    def forward_features(self, x):
-        #x, sizes = self.forward_network(x)
         size: Tuple[int, int] = (x.size(2), x.size(3))
-        x, size = self.stages(x, size)
+        features = [x]
+        sizes = [size]
+        
+        for stage in self.stages:
+            features[-1], sizes[-1] = stage(features[-1], sizes[-1])
+            
+            # don't append outputs of last stage, since they are already there
+            if(len(features) < self.num_stages):
+                features.append(features[-1])
+                sizes.append(sizes[-1])
+            
+
+        # non-normalized pyramid features + corresponding sizes
+        return features, sizes
     
+    def forward_pyramid_features(self, x) -> List[Tensor]:
+        x, sizes = self.forward_network(x)
+        outs = []
+        for i, out in enumerate(x):
+            H, W = sizes[i]
+            outs.append(out.view(-1, H, W, self.embed_dims[i]).permute(0, 3, 1, 2).contiguous())       
+        return outs
+
+    def forward_features(self, x):
+        x, sizes = self.forward_network(x)
+        # take final feature and norm
         x = self.norms(x)
         H, W = size
         x = x.view(-1, H, W, self.embed_dims[-1]).permute(0, 3, 1, 2).contiguous()
@@ -570,6 +592,14 @@ class DaViT(nn.Module):
     def forward(self, x):
         return self.forward_classifier(x)
 
+class DaViTFeatures(DaViT):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.feature_info = FeatureInfo(self.feature_info, kwargs.get('out_indices', (0, 1, 2, 3)))
+    
+    def forward(self, x) -> List[Tensor]:
+        return self.forward_pyramid_features(x)
 
 def checkpoint_filter_fn(state_dict, model):
     """ Remap MSFT checkpoints -> timm """
@@ -593,18 +623,30 @@ def checkpoint_filter_fn(state_dict, model):
 
 def _create_davit(variant, pretrained=False, **kwargs):
 
+    model_cls = DaViT
+    features_only = False
+    kwargs_filter = None
+
     default_out_indices = tuple(i for i, _ in enumerate(kwargs.get('depths', (1, 1, 3, 1))))
     out_indices = kwargs.pop('out_indices', default_out_indices)
+    if kwargs.pop('features_only', False):
+        model_cls = DaViTFeatures
+        kwargs_filter = ('num_classes', 'global_pool')
+        features_only = True
     model = build_model_with_cfg(
-        DaViT,
+        model_cls,
         variant,
         pretrained,
         pretrained_filter_fn=checkpoint_filter_fn,
         feature_cfg=dict(flatten_sequential=True, out_indices=out_indices),
         **kwargs)
+    if features_only:
+        model.pretrained_cfg = pretrained_cfg_for_features(model.default_cfg)
+        model.default_cfg = model.pretrained_cfg  # backwards compat
     return model
     
     
+
 def _cfg(url='', **kwargs): # not sure how this should be set up
     return {
         'url': url,