fix some model

Author: brianhou0208

timm/models/byobnet.py

@@ -44,7 +44,7 @@
 )
 from ._builder import build_model_with_cfg
 from ._features import feature_take_indices
-from ._manipulate import checkpoint, checkpoint_seq, named_apply
+from ._manipulate import named_apply, checkpoint_seq
 from ._registry import generate_default_cfgs, register_model
 
 __all__ = ['ByobNet', 'ByoModelCfg', 'ByoBlockCfg', 'create_byob_stem', 'create_block']
@@ -1385,7 +1385,7 @@ def forward_intermediates(
         for stage in stages:
             feat_idx += 1
             if self.grad_checkpointing and not torch.jit.is_scripting():
-                x = checkpoint(stage, x)
+                x = checkpoint_seq(stage, x)
             else:
                 x = stage(x)
             if not exclude_final_conv and feat_idx == last_idx:

timm/models/efficientnet.py

@@ -212,7 +212,7 @@ def forward_intermediates(
             blocks = self.blocks[:max_index]
         for feat_idx, blk in enumerate(blocks, start=1):
             if self.grad_checkpointing and not torch.jit.is_scripting():
-                x = checkpoint(blk, x)
+                x = checkpoint_seq(blk, x, flatten=True)
             else:
                 x = blk(x)
             if feat_idx in take_indices:

timm/models/efficientvit_mit.py

@@ -790,7 +790,7 @@ def forward_intermediates(
 
         for feat_idx, stage in enumerate(stages):
             if self.grad_checkpointing and not torch.jit.is_scripting():
-                x = checkpoint(stage, x)
+                x = checkpoint_seq(stages, x)
             else:
                 x = stage(x)
             if feat_idx in take_indices:
@@ -947,7 +947,7 @@ def forward_intermediates(
 
         for feat_idx, stage in enumerate(stages):
             if self.grad_checkpointing and not torch.jit.is_scripting():
-                x = checkpoint(stage, x)
+                x = checkpoint_seq(stages, x)
             else:
                 x = stage(x)
             if feat_idx in take_indices:

timm/models/ghostnet.py

@@ -728,7 +728,7 @@ def forward_intermediates(
 
         for feat_idx, stage in enumerate(stages, start=1):
             if self.grad_checkpointing and not torch.jit.is_scripting():
-                x = checkpoint(stage, x)
+                x = checkpoint_seq(stage, x, flatten=True)
             else:
                 x = stage(x)
             if feat_idx in take_indices:

timm/models/hieradet_sam2.py

@@ -1,21 +1,20 @@
 import math
 from copy import deepcopy
 from functools import partial
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.jit import Final
 
 from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
 from timm.layers import PatchEmbed, Mlp, DropPath, ClNormMlpClassifierHead, LayerScale, \
     get_norm_layer, get_act_layer, init_weight_jax, init_weight_vit, to_2tuple, use_fused_attn
 
 from ._builder import build_model_with_cfg
 from ._features import feature_take_indices
-from ._manipulate import named_apply, checkpoint_seq, adapt_input_conv
-from ._registry import generate_default_cfgs, register_model, register_model_deprecations
+from ._manipulate import named_apply, checkpoint
+from ._registry import generate_default_cfgs, register_model
 
 
 def window_partition(x, window_size: Tuple[int, int]):
@@ -471,7 +470,10 @@ def forward_intermediates(
         else:
             blocks = self.blocks[:max_index + 1]
         for i, blk in enumerate(blocks):
-            x = blk(x)
+            if self.grad_checkpointing and not torch.jit.is_scripting():
+                x = checkpoint(blk, x)
+            else:
+                x = blk(x)
             if i in take_indices:
                 x_out = x.permute(0, 3, 1, 2) if output_fmt == 'NCHW' else x
                 intermediates.append(x_out)
@@ -503,8 +505,11 @@ def prune_intermediate_layers(
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.patch_embed(x)  # BHWC
         x = self._pos_embed(x)
-        for i, blk in enumerate(self.blocks):
-            x = blk(x)
+        for blk in self.blocks:
+            if self.grad_checkpointing and not torch.jit.is_scripting():
+                x = checkpoint(blk, x)
+            else:
+                x = blk(x)
         return x
 
     def forward_head(self, x, pre_logits: bool = False) -> torch.Tensor:

timm/models/inception_resnet_v2.py

@@ -5,7 +5,6 @@
 from functools import partial
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 from timm.data import IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD
 from timm.layers import create_classifier, ConvNormAct

timm/models/inception_v3.py

@@ -4,7 +4,6 @@
 Licensed BSD-Clause 3 https://github.com/pytorch/vision/blob/master/LICENSE
 """
 from functools import partial
-from typing import Optional
 
 import torch
 import torch.nn as nn

timm/models/nasnet.py

@@ -3,11 +3,9 @@
  https://github.com/Cadene/pretrained-models.pytorch
 """
 from functools import partial
-from typing import Optional
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 from timm.layers import ConvNormAct, create_conv2d, create_pool2d, create_classifier
 from ._builder import build_model_with_cfg

timm/models/nextvit.py

@@ -17,7 +17,7 @@
 from timm.layers import ClassifierHead
 from ._builder import build_model_with_cfg
 from ._features import feature_take_indices
-from ._manipulate import checkpoint, checkpoint_seq
+from ._manipulate import checkpoint_seq
 from ._registry import generate_default_cfgs, register_model
 
 __all__ = ['NextViT']
@@ -595,7 +595,7 @@ def forward_intermediates(
 
         for feat_idx, stage in enumerate(stages):
             if self.grad_checkpointing and not torch.jit.is_scripting():
-                x = checkpoint(stage, x)
+                x = checkpoint_seq(stage, x)
             else:
                 x = stage(x)
             if feat_idx in take_indices:

timm/models/nfnet.py

@@ -215,7 +215,7 @@ def create_stem(
     if 'deep' in stem_type:
         if 'quad' in stem_type:
             # 4 deep conv stack as in NFNet-F models
-            assert not 'pool' in stem_type
+            assert 'pool' not in stem_type
             stem_chs = (out_chs // 8, out_chs // 4, out_chs // 2, out_chs)
             strides = (2, 1, 1, 2)
             stem_stride = 4