Fix torch version comparison (#10934)

.dev_scripts/gather_models.py

@@ -12,7 +12,7 @@
 import yaml
 from mmengine.config import Config
 from mmengine.fileio import dump
-from mmengine.utils import mkdir_or_exist, scandir
+from mmengine.utils import digit_version, mkdir_or_exist, scandir
 
 
 def ordered_yaml_dump(data, stream=None, Dumper=yaml.SafeDumper, **kwds):
@@ -45,7 +45,7 @@ def process_checkpoint(in_file, out_file):
 
     # if it is necessary to remove some sensitive data in checkpoint['meta'],
     # add the code here.
-    if torch.__version__ >= '1.6':
+    if digit_version(torch.__version__) >= digit_version('1.6'):
         torch.save(checkpoint, out_file, _use_new_zipfile_serialization=False)
     else:
         torch.save(checkpoint, out_file)

mmdet/models/layers/normed_predictor.py

@@ -2,6 +2,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from mmengine.utils import digit_version
 from torch import Tensor
 
 from mmdet.registry import MODELS
@@ -91,7 +92,7 @@ def forward(self, x: Tensor) -> Tensor:
         if hasattr(self, 'conv2d_forward'):
             x_ = self.conv2d_forward(x_, weight_)
         else:
-            if torch.__version__ >= '1.8':
+            if digit_version(torch.__version__) >= digit_version('1.8'):
                 x_ = self._conv_forward(x_, weight_, self.bias)
             else:
                 x_ = self._conv_forward(x_, weight_)

tools/analysis_tools/get_flops.py

@@ -11,6 +11,7 @@
 from mmengine.model import revert_sync_batchnorm
 from mmengine.registry import init_default_scope
 from mmengine.runner import Runner
+from mmengine.utils import digit_version
 
 from mmdet.registry import MODELS
 
@@ -44,7 +45,7 @@ def parse_args():
 
 
 def inference(args, logger):
-    if str(torch.__version__) < '1.12':
+    if digit_version(torch.__version__) < digit_version('1.12'):
         logger.warning(
             'Some config files, such as configs/yolact and configs/detectors,'
             'may have compatibility issues with torch.jit when torch<1.12. '

tools/model_converters/publish_model.py

@@ -4,6 +4,7 @@
 
 import torch
 from mmengine.logging import print_log
+from mmengine.utils import digit_version
 
 
 def parse_args():
@@ -37,7 +38,7 @@ def process_checkpoint(in_file, out_file, save_keys=['meta', 'state_dict']):
 
     # if it is necessary to remove some sensitive data in checkpoint['meta'],
     # add the code here.
-    if torch.__version__ >= '1.6':
+    if digit_version(torch.__version__) >= digit_version('1.6'):
         torch.save(checkpoint, out_file, _use_new_zipfile_serialization=False)
     else:
         torch.save(checkpoint, out_file)

tools/model_converters/upgrade_ssd_version.py

@@ -5,6 +5,7 @@
 
 import torch
 from mmengine import Config
+from mmengine.utils import digit_version
 
 
 def parse_config(config_strings):
@@ -39,7 +40,7 @@ def convert(in_file, out_file):
         out_state_dict[new_key] = value
     checkpoint['state_dict'] = out_state_dict
 
-    if torch.__version__ >= '1.6':
+    if digit_version(torch.__version__) >= digit_version('1.6'):
         torch.save(checkpoint, out_file, _use_new_zipfile_serialization=False)
     else:
         torch.save(checkpoint, out_file)