Fix early stop for efficientnet/mobilenetv3 fwd inter. Fix indices typing for all fwd inter.

Author: rwightman

timm/models/convnext.py

@@ -411,7 +411,7 @@ def reset_classifier(self, num_classes=0, global_pool=None):
     def forward_intermediates(
             self,
             x: torch.Tensor,
-            indices: Union[int, List[int], Tuple[int]] = None,
+            indices: Optional[Union[int, List[int], Tuple[int]]] = None,
             norm: bool = False,
             stop_early: bool = False,
             output_fmt: str = 'NCHW',

timm/models/efficientformer.py

@@ -12,7 +12,7 @@
 
 Modifications and timm support by / Copyright 2022, Ross Wightman
 """
-from typing import Dict, List, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -463,7 +463,7 @@ def set_distilled_training(self, enable=True):
     def forward_intermediates(
             self,
             x: torch.Tensor,
-            indices: Union[int, List[int], Tuple[int]] = None,
+            indices: Optional[Union[int, List[int], Tuple[int]]] = None,
             norm: bool = False,
             stop_early: bool = False,
             output_fmt: str = 'NCHW',

timm/models/efficientnet.py

@@ -162,7 +162,7 @@ def reset_classifier(self, num_classes, global_pool='avg'):
     def forward_intermediates(
             self,
             x: torch.Tensor,
-            indices: Union[int, List[int], Tuple[int]] = None,
+            indices: Optional[Union[int, List[int], Tuple[int]]] = None,
             norm: bool = False,
             stop_early: bool = False,
             output_fmt: str = 'NCHW',
@@ -183,8 +183,6 @@ def forward_intermediates(
 
         """
         assert output_fmt in ('NCHW',), 'Output shape must be NCHW.'
-        if stop_early:
-            assert intermediates_only, 'Must use intermediates_only for early stopping.'
         intermediates = []
         if extra_blocks:
             take_indices, max_index = feature_take_indices(len(self.blocks) + 1, indices)
@@ -212,8 +210,9 @@ def forward_intermediates(
         if intermediates_only:
             return intermediates
 
-        x = self.conv_head(x)
-        x = self.bn2(x)
+        if feat_idx == self.stage_ends[-1]:
+            x = self.conv_head(x)
+            x = self.bn2(x)
 
         return x, intermediates
 

timm/models/levit.py

@@ -25,7 +25,7 @@
 # Copyright 2020 Ross Wightman, Apache-2.0 License
 from collections import OrderedDict
 from functools import partial
-from typing import Dict, List, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -638,7 +638,7 @@ def reset_classifier(self, num_classes, global_pool=None, distillation=None):
     def forward_intermediates(
             self,
             x: torch.Tensor,
-            indices: Union[int, List[int], Tuple[int]] = None,
+            indices: Optional[Union[int, List[int], Tuple[int]]] = None,
             norm: bool = False,
             stop_early: bool = False,
             output_fmt: str = 'NCHW',

timm/models/maxxvit.py

@@ -1255,7 +1255,7 @@ def reset_classifier(self, num_classes, global_pool=None):
     def forward_intermediates(
             self,
             x: torch.Tensor,
-            indices: Union[int, List[int], Tuple[int]] = None,
+            indices: Optional[Union[int, List[int], Tuple[int]]] = None,
             norm: bool = False,
             stop_early: bool = False,
             output_fmt: str = 'NCHW',

timm/models/mobilenetv3.py

@@ -154,7 +154,7 @@ def reset_classifier(self, num_classes: int, global_pool: str = 'avg'):
     def forward_intermediates(
             self,
             x: torch.Tensor,
-            indices: Union[int, List[int], Tuple[int]] = None,
+            indices: Optional[Union[int, List[int], Tuple[int]]] = None,
             norm: bool = False,
             stop_early: bool = False,
             output_fmt: str = 'NCHW',

timm/models/mvitv2.py

@@ -837,7 +837,7 @@ def reset_classifier(self, num_classes, global_pool=None):
     def forward_intermediates(
             self,
             x: torch.Tensor,
-            indices: Union[int, List[int], Tuple[int]] = None,
+            indices: Optional[Union[int, List[int], Tuple[int]]] = None,
             norm: bool = False,
             stop_early: bool = False,
             output_fmt: str = 'NCHW',

timm/models/resnet.py

@@ -557,7 +557,7 @@ def reset_classifier(self, num_classes, global_pool='avg'):
     def forward_intermediates(
             self,
             x: torch.Tensor,
-            indices: Union[int, List[int], Tuple[int]] = None,
+            indices: Optional[Union[int, List[int], Tuple[int]]] = None,
             norm: bool = False,
             stop_early: bool = False,
             output_fmt: str = 'NCHW',
@@ -576,8 +576,6 @@ def forward_intermediates(
 
         """
         assert output_fmt in ('NCHW',), 'Output shape must be NCHW.'
-        if stop_early:
-            assert intermediates_only, 'Must use intermediates_only for early stopping.'
         intermediates = []
         take_indices, max_index = feature_take_indices(5, indices)
 

timm/models/swin_transformer.py

@@ -611,7 +611,7 @@ def reset_classifier(self, num_classes, global_pool=None):
     def forward_intermediates(
             self,
             x: torch.Tensor,
-            indices: Union[int, List[int], Tuple[int]] = None,
+            indices: Optional[Union[int, List[int], Tuple[int]]] = None,
             norm: bool = False,
             stop_early: bool = False,
             output_fmt: str = 'NCHW',

timm/models/swin_transformer_v2.py

@@ -612,7 +612,7 @@ def reset_classifier(self, num_classes, global_pool=None):
     def forward_intermediates(
             self,
             x: torch.Tensor,
-            indices: Union[int, List[int], Tuple[int]] = None,
+            indices: Optional[Union[int, List[int], Tuple[int]]] = None,
             norm: bool = False,
             stop_early: bool = False,
             output_fmt: str = 'NCHW',