add tuning code

Author: lyuwenyu

README.md

@@ -36,6 +36,7 @@ This is the official implementation of the paper "[DETRs Beat YOLOs on Real-time
 
 ## Updates!!!
 ---
+- \[2023.10.12\] Add tuning code for pytorch version, now you can tuning rtdetr based on pretrained weights
 - \[2023.09.19\] Upload [*pytorch weights*](https://github.com/lyuwenyu/RT-DETR/issues/42) convert from paddle version
 - \[2023.08.24] Release rtdetr-18 pretrained models on objects365. *49.2 mAP* and *217 FPS*
 - \[2023.08.22\] Upload *[rtdetr_pytorch](./rtdetr_pytorch/)* source code. Please enjoy it ❤️

rtdetr_pytorch/README.md

@@ -5,7 +5,7 @@
 - [x] Upload source code
 - [x] Upload weight convert from paddle, see [links](https://github.com/lyuwenyu/RT-DETR/issues/42)
 - [x] Align training details with the [paddle version](../rtdetr_paddle/)
-
+- [x] Tuning rtdetr based on [pretrained weights](https://github.com/lyuwenyu/RT-DETR/issues/42)
 
 ## Quick start
 
@@ -79,6 +79,7 @@ python tools/export_onnx.py -c configs/rtdetr/rtdetr_r18vd_6x_coco.yml -r path/t
 <details open>
 <summary>Train custom data</summary>
 
-set `remap_mscoco_category: False`. This variable only works for ms-coco dataset.
+1. set `remap_mscoco_category: False`. This variable only works for ms-coco dataset.
 
+2. add `-t path/to/checkpoint` (optinal) to tuning rtdetr based on pretrained checkpoint. see [training script details](./tools/README.md).
 </details>

rtdetr_pytorch/src/core/config.py

@@ -50,6 +50,7 @@ def __init__(self) -> None:
 
         # runtime
         self.resume :str = None
+        self.tuning :str = None
 
         self.epoches :int = None
         self.last_epoch :int = -1

rtdetr_pytorch/src/core/yaml_config.py

@@ -26,6 +26,7 @@ def __init__(self, cfg_path: str, **kwargs) -> None:
         self.checkpoint_step = cfg.get('checkpoint_step', 1)
         self.epoches = cfg.get('epoches', -1)
         self.resume = cfg.get('resume', '')
+        self.tuning = cfg.get('tuning', '')
         self.sync_bn = cfg.get('sync_bn', False)
         self.output_dir = cfg.get('output_dir', None)
 

rtdetr_pytorch/src/misc/dist.py

@@ -20,7 +20,7 @@
 from torch.utils.data.dataloader import DataLoader
 
 
-def init_distributed(backend='nccl'):
+def init_distributed():
     '''
     distributed setup
     args:
@@ -32,7 +32,7 @@ def init_distributed(backend='nccl'):
         # RANK = int(os.getenv('RANK', -1))
         # WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
 
-        tdist.init_process_group(backend=backend, init_method='env://', )
+        tdist.init_process_group(init_method='env://', )
         torch.distributed.barrier()
 
         rank = get_rank()

rtdetr_pytorch/src/solver/solver.py

@@ -6,6 +6,7 @@
 
 from datetime import datetime
 from pathlib import Path 
+from typing import Dict
 
 from src.misc import dist
 from src.core import BaseConfig
@@ -28,6 +29,11 @@ def setup(self, ):
         self.criterion = cfg.criterion.to(device)
         self.postprocessor = cfg.postprocessor
 
+        # NOTE (lvwenyu): should load_tuning_state before ema instance building
+        if self.cfg.tuning:
+            print(f'Tuning checkpoint from {self.cfg.tuning}')
+            self.load_tuning_state(self.cfg.tuning)
+
         self.scaler = cfg.scaler
         self.ema = cfg.ema.to(device) if cfg.ema is not None else None 
 
@@ -133,10 +139,44 @@ def resume(self, path):
         state = torch.load(path, map_location='cpu')
         self.load_state_dict(state)
 
+    def load_tuning_state(self, path,):
+        """only load model for tuning and skip missed/dismatched keys
+        """
+        if 'http' in path:
+            state = torch.hub.load_state_dict_from_url(path, map_location='cpu')
+        else:
+            state = torch.load(path, map_location='cpu')
+
+        module = dist.de_parallel(self.model)
+        
+        # TODO hard code
+        if 'ema' in state:
+            stat, infos = self._matched_state(module.state_dict(), state['ema']['module'])
+        else:
+            stat, infos = self._matched_state(module.state_dict(), state['model'])
+
+        module.load_state_dict(stat, strict=False)
+        print(f'Load model.state_dict, {infos}')
+
+    @staticmethod
+    def _matched_state(state: Dict[str, torch.Tensor], params: Dict[str, torch.Tensor]):
+        missed_list = []
+        unmatched_list = []
+        matched_state = {}
+        for k, v in state.items():
+            if k in params:
+                if v.shape == params[k].shape:
+                    matched_state[k] = params[k]
+                else:
+                    unmatched_list.append(k)
+            else:
+                missed_list.append(k)
+
+        return matched_state, {'missed': missed_list, 'unmatched': unmatched_list}
+
 
     def fit(self, ):
         raise NotImplementedError('')
 
-
     def val(self, ):
         raise NotImplementedError('')

rtdetr_pytorch/tools/README.md

@@ -7,6 +7,10 @@ Train/test script examples
 - `--test-only` 
 
 
+Tuning script examples
+- `torchrun --master_port=8844 --nproc_per_node=4 tools/train.py -c configs/rtdetr/rtdetr_r18vd_6x_coco.yml -t https://github.com/lyuwenyu/storage/releases/download/v0.1/rtdetr_r18vd_5x_coco_objects365_from_paddle.pth` 
+
+
 Export script examples
 - `python tools/export_onnx.py -c path/to/config -r path/to/checkpoint --check`
 

rtdetr_pytorch/tools/train.py

@@ -14,8 +14,18 @@
 def main(args, ) -> None:
     '''main
     '''
-    dist.init_distributed(backend='nccl')
-    cfg = YAMLConfig(args.config, resume=args.resume, use_amp=args.amp)
+    dist.init_distributed()
+
+    assert not all([args.tuning, args.resume]), \
+        'Only support from_scrach or resume or tuning at one time'
+
+    cfg = YAMLConfig(
+        args.config,
+        resume=args.resume, 
+        use_amp=args.amp,
+        tuning=args.tuning
+    )
+
     solver = TASKS[cfg.yaml_cfg['task']](cfg)
 
     if args.test_only:
@@ -24,12 +34,12 @@ def main(args, ) -> None:
         solver.fit()
 
 
-
 if __name__ == '__main__':
 
     parser = argparse.ArgumentParser()
     parser.add_argument('--config', '-c', type=str, )
     parser.add_argument('--resume', '-r', type=str, )
+    parser.add_argument('--tuning', '-t', type=str, )
     parser.add_argument('--test-only', action='store_true', default=False,)
     parser.add_argument('--amp', action='store_true', default=False,)