[update] 20260312

Author: Lmy0217

.github/workflows/probot.yml

@@ -8,7 +8,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.9', '3.10', '3.11', '3.12']
+        python-version: ['3.10', '3.11', '3.12', '3.13']
 
     steps:
     - uses: actions/checkout@v1
@@ -19,8 +19,8 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install -r requirements.txt
         pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cpu
+        pip install -r requirements.txt
     - name: Lint with flake8
       run: |
         pip install flake8

.travis.yml

@@ -6,10 +6,12 @@ python:
   - "3.10"
   - "3.11"
   - "3.12"
+  - "3.13"
 
 install:
-  - pip install -r requirements.txt
+  - python -m pip install --upgrade pip
   - pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cpu
+  - pip install -r requirements.txt
 
 script:
   - python -m unittest discover test

circle.yml

@@ -3,16 +3,16 @@ version: 2.1
 workflows:
   test:
     jobs:
-      - job-py38
       - job-py39
       - job-py310
       - job-py311
       - job-py312
+      - job-py313
 
 jobs:
-  job-py38: &job-template
+  job-py39: &job-template
     docker:
-      - image: cimg/python:3.8
+      - image: cimg/python:3.9
 
     working_directory: ~/repo
 
@@ -25,8 +25,8 @@ jobs:
             python -m venv venv
             . venv/bin/activate
             pip install --upgrade pip
-            pip install -r requirements.txt > /dev/null
             pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cpu
+            pip install -r requirements.txt
 
       - run:
           name: run unittest
@@ -50,11 +50,6 @@ jobs:
           path: test-reports
           destination: test-reports
 
-  job-py39:
-    <<: *job-template
-    docker:
-      - image: cimg/python:3.9
-
   job-py310:
     <<: *job-template
     docker:
@@ -69,3 +64,8 @@ jobs:
     <<: *job-template
     docker:
       - image: cimg/python:3.12
+
+  job-py313:
+    <<: *job-template
+    docker:
+      - image: cimg/python:3.13

configs/BaseConfig.py

@@ -25,6 +25,8 @@ def dict(self):
         for key in list(value_dict.keys()):
             if not self._values(key):
                 value_dict.pop(key, None)
+            elif isinstance(value_dict[key], BaseConfig):
+                value_dict[key] = value_dict[key].dict()
         return value_dict
 
     def __eq__(self, other):

datasets/BaseDataset.py

@@ -18,6 +18,7 @@
 class BaseDataset(Dataset, metaclass=abc.ABCMeta):
 
     logger: utils.Logger
+    wandb: utils.Wandb
     summary: utils.Summary
 
     def __init__(self, cfg, **kwargs):
@@ -55,8 +56,15 @@ def set_logger(self, logger):
         if hasattr(self, 'super_dataset') and isinstance(self.super_dataset, BaseDataset):
             self.super_dataset.set_logger(logger)
 
+    def set_wandb(self, wandb):
+        self.wandb = wandb
+        if hasattr(self, 'super_dataset') and isinstance(self.super_dataset, BaseDataset):
+            self.super_dataset.set_wandb(wandb)
+
     def set_summary(self, summary):
         self.summary = summary
+        if hasattr(self, 'super_dataset') and isinstance(self.super_dataset, BaseDataset):
+            self.super_dataset.set_summary(summary)
 
     @staticmethod
     def need_norm(data_shape):
@@ -153,17 +161,17 @@ def _get_args(self, norm_func, data, name: str, **kwargs):
     def _recover(self, index):
         index_dict = dict()
         for key, value in self.cfg.dict().items():
-            if key not in ['out', 'kernel'] and isinstance(value, configs.BaseConfig):
-                if hasattr(value, 'width') and hasattr(value, 'height'):
-                    if hasattr(value, 'time'):
-                        if hasattr(value, 'patch'):
-                            index_dict[key] = [0, value.patch, 0, value.time, 0, value.width, 0, value.height]
+            if key not in ['out', 'kernel'] and isinstance(value, dict):
+                if 'width' in value and 'height' in value:
+                    if 'time' in value:
+                        if 'patch' in value:
+                            index_dict[key] = [0, value['patch'], 0, value['time'], 0, value['width'], 0, value['height']]
                         else:
-                            index_dict[key] = [0, value.time, 0, value.width, 0, value.height]
+                            index_dict[key] = [0, value['time'], 0, value['width'], 0, value['height']]
                     else:
-                        index_dict[key] = [0, value.width, 0, value.height]
-                elif hasattr(value, 'elements'):
-                    index_dict[key] = [0, value.elements]
+                        index_dict[key] = [0, value['width'], 0, value['height']]
+                elif 'elements' in value:
+                    index_dict[key] = [0, value['elements']]
         return [index], index_dict
 
     def recover(self, index):

main.py

@@ -51,6 +51,8 @@ def split(self, index_cross):
 
         self.logger = utils.Logger(self.path, utils.path.get_filename(self.model_cfg._path))
         self.dataset.set_logger(self.logger)
+        self.wandb = utils.Wandb(self.model_cfg, self.dataset_cfg, self.run_cfg)
+        self.dataset.set_wandb(self.wandb)
         self.summary = utils.Summary(self.path, dataset=self.dataset)
         self.dataset.set_summary(self.summary)
 
@@ -97,7 +99,8 @@ def split(self, index_cross):
         )
 
         self.model = models.functional.common.find(self.model_cfg.name)(
-            self.model_cfg, self.dataset.cfg, self.run_cfg, logger=self.logger, summary=self.summary, main_msg=self.msg)
+            self.model_cfg, self.dataset.cfg, self.run_cfg,
+            logger=self.logger, wandb=self.wandb, summary=self.summary, main_msg=self.msg)
         self.start_epoch = self.model.load(self.args.test_epoch)
 
         if not self.run_cfg.distributed or (self.run_cfg.distributed and self.run_cfg.local_rank == 0):
@@ -121,7 +124,7 @@ def train(self, epoch):
         if self.run_cfg.distributed:
             self.train_loader.sampler.set_epoch(epoch)
         self.train_loader = self.model.train_loader_hook(self.train_loader)
-        batch_per_epoch, count_data = len(self.train_loader), len(self.train_loader.dataset)
+        batch_per_epoch, count_data = len(self.train_loader), len(self.train_loader.sampler)
         log_step = 1#max(int(np.power(10, np.floor(np.log10(batch_per_epoch / 10)))), 1) if batch_per_epoch > 0 else 1
         epoch_info = {'epoch': epoch, 'batch_per_epoch': batch_per_epoch, 'count_data': count_data}
         epoch_info['local_rank'] = self.run_cfg.local_rank if self.run_cfg.distributed else 1
@@ -165,8 +168,10 @@ def train(self, epoch):
         if self.run_cfg.distributed:
             with utils.ddp.sequence():
                 self.logger.info_scalars('Train Epoch: {} rank {}\t', (epoch, self.run_cfg.local_rank), loss_all)
+                self.wandb.info(loss_all)
         else:
             self.logger.info_scalars('Train Epoch: {}\t', (epoch,), loss_all)
+            self.wandb.info(loss_all)
         if epoch % self.run_cfg.save_step == 0:
             self.model.save(epoch)
 
@@ -182,7 +187,7 @@ def test(self, epoch, data_loader=None, log_text='Test'):
             if self.run_cfg.distributed:
                 data_loader.sampler.set_epoch(epoch)
             data_loader = self.model.test_loader_hook(data_loader)
-            batch_per_epoch, count_data = len(data_loader), len(data_loader.dataset)
+            batch_per_epoch, count_data = len(data_loader), len(data_loader.sampler)
             log_step = max(int(np.power(10, np.floor(np.log10(batch_per_epoch / 10)))), 1)
             epoch_info = {'epoch': epoch, 'batch_per_epoch': batch_per_epoch, 'count_data': count_data, 'log_text': log_text}
             epoch_info['local_rank'] = self.run_cfg.local_rank if self.run_cfg.distributed else 1
@@ -269,6 +274,7 @@ def test(self, epoch, data_loader=None, log_text='Test'):
                     msgs_dict[name] = 100. * value / count
                     accuracy.append(msgs_dict[name])
                 self.logger.info(log_msg.format(log_text, epoch, *accuracy))
+                self.wandb.info({f'{log_text}_accuracy': accuracy})
                 self.summary.add_scalars('Accuracy', msgs_dict, epoch)
 
         # TODO do not support chain norm and renorm
@@ -277,7 +283,7 @@ def test(self, epoch, data_loader=None, log_text='Test'):
         #     dataset.append(dataset[-1].super_dataset)
         #     dataset_cfg.append(dataset[-1].cfg)
         for name, value in predict.items():
-            predict[name] = np.array(value.cpu())
+            predict[name] = value.cpu().numpy()
             for d_cfg, d in zip(dataset_cfg, dataset):
                 if d_cfg.norm and d.need_norm(value.shape):
                     data_type, data_cfg = self._get_type(d_cfg, name, test=False)
@@ -339,17 +345,22 @@ def run():
             main.split(index_cross)
             main.model.process_pre_hook()
             if args.test_epoch is None:
+                main.model.main_msg.update(dict(only_test=False))
                 if main.start_epoch == 0:
-                    main.val_test(main.start_epoch)  # TODO set flag only_test=False
+                    main.model.process_test_msg_pre_hook(main.model.main_msg)
+                    while main.model.main_msg['test_flag']:
+                        main.val_test(main.start_epoch)
+                        main.model.process_test_msg_hook(main.model.main_msg)
                 for epoch in range(main.start_epoch + 1, main.run_cfg.epochs + 1):
                     main.train(epoch)
                     if epoch % main.run_cfg.save_step == 0:
-                        main.model.main_msg.update(dict(test_idx=1, test_flag=True, only_test=False))
+                        main.model.process_test_msg_pre_hook(main.model.main_msg)
                         while main.model.main_msg['test_flag']:
                             main.val_test(epoch)
                             main.model.process_test_msg_hook(main.model.main_msg)
             else:
-                main.model.main_msg.update(dict(test_idx=1, test_flag=True, only_test=True))
+                main.model.main_msg.update(dict(only_test=True))
+                main.model.process_test_msg_pre_hook(main.model.main_msg)
                 while main.model.main_msg['test_flag']:
                     main.val_test(main.start_epoch)
                     main.model.process_test_msg_hook(main.model.main_msg)

models/BaseModel.py

@@ -3,6 +3,8 @@
 
 import torch
 import torch.nn as nn
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import DataLoader
 
 import configs
@@ -23,6 +25,9 @@ def process_hook(self):
     def process_msg_hook(self, msg: dict):
         msg.update(dict(while_flag=False))
 
+    def process_test_msg_pre_hook(self, msg: dict):
+        msg.update(dict(test_flag=True, test_idx=1))
+
     def process_test_msg_hook(self, msg: dict):
         msg.update(dict(test_flag=False))
 
@@ -83,6 +88,7 @@ def test_return_hook(self, epoch_info: dict, return_all: dict):
 class BaseModel(_ProcessHook, _MainHook, metaclass=abc.ABCMeta):
 
     logger: utils.Logger
+    wandb: utils.Wandb
     summary: utils.Summary
     main_msg: dict
 
@@ -145,7 +151,7 @@ def summary_models(self, shapes):
     def load(self, start_epoch=None, path=None):
         assert start_epoch is None or (isinstance(start_epoch, int) and start_epoch >= 0)
         path = path or self.path
-        load_optimizer = start_epoch is None
+        load_optimizer_scheduler = start_epoch is None
         if start_epoch is None:
             main_msg = ('_' + str(self.main_msg['while_idx'])) if self.main_msg['while_idx'] > 1 else ''
             check_path = os.path.join(path, self.name + main_msg + configs.env.paths.check_file)
@@ -159,16 +165,24 @@ def load(self, start_epoch=None, path=None):
         if start_epoch > 0:
             msg = ('_' + '-'.join(self.msg.values())) if self.msg else ''
             for name, value in self.__dict__.items():
-                if isinstance(value, (nn.Module, torch.optim.Optimizer)) or name in self._save_list:
-                    epoch_msg = ('_' + str(start_epoch)) if not isinstance(value, torch.optim.Optimizer) else ''
+                if isinstance(value, (nn.Module, Optimizer, LRScheduler)) or name in self._save_list:
+                    epoch_msg = ('_' + str(start_epoch)) if not isinstance(value, (Optimizer, LRScheduler)) else ''
                     load_path = os.path.join(path, self.name + '_' + name + epoch_msg + msg + '.pth')
-                    if not (load_optimizer and os.path.exists(load_path)) and isinstance(value, torch.optim.Optimizer):
-                        if not os.path.exists(load_path):
-                            self.logger.info(f"IGNORE! Optimizer weight `{load_path}` not found!")
+                    if not load_optimizer_scheduler and isinstance(value, (Optimizer, LRScheduler)):
                         continue
+                    if not os.path.exists(load_path):
+                        if isinstance(value, Optimizer):
+                            self.logger.info(f"IGNORE! Optimizer `{load_path}` not found!")
+                            continue
+                        elif isinstance(value, LRScheduler):
+                            self.logger.info(f"IGNORE! LRScheduler `{load_path}` not found!")
+                            for _ in range(start_epoch):
+                                value.step()
+                            self.logger.info(f"LRScheduler `step()` has been executed automatically {start_epoch} times!")
+                            continue
                     map_location = {'cuda:%d' % 0: 'cuda:%d' % self.run.local_rank} if self.run.distributed else self.device
                     load_value = torch.load(load_path, map_location=map_location)
-                    if isinstance(value, (nn.Module, torch.optim.Optimizer)):
+                    if isinstance(value, (nn.Module, Optimizer, LRScheduler)):
                         self.__dict__[name].load_state_dict(load_value)
                     else:
                         self.__dict__[name] = load_value
@@ -182,9 +196,9 @@ def save(self, epoch, path=None):
             msg = ('_' + '-'.join(self.msg.values())) if self.msg else ''
             for name, value in self.__dict__.items():
                 # TODO remove criterion, change criterion super object to `torch.nn.modules.loss._Loss`?
-                if isinstance(value, (nn.Module, torch.optim.Optimizer)) or name in self._save_list:
-                    save_value = value.state_dict() if isinstance(value, (nn.Module, torch.optim.Optimizer)) else value
-                    epoch_msg = ('_' + str(epoch)) if not isinstance(value, torch.optim.Optimizer) else ''
+                if isinstance(value, (nn.Module, Optimizer, LRScheduler)) or name in self._save_list:
+                    save_value = value.state_dict() if isinstance(value, (nn.Module, Optimizer, LRScheduler)) else value
+                    epoch_msg = ('_' + str(epoch)) if not isinstance(value, (Optimizer, LRScheduler)) else ''
                     torch.save(save_value, os.path.join(path, self.name + '_' + name + epoch_msg + msg + '.pth'))
             main_msg = ('_' + str(self.main_msg['while_idx'])) if self.main_msg['while_idx'] > 1 else ''
             torch.save(dict(epoch=epoch, msg=self.msg, main_msg=self.main_msg),

requirements.txt

@@ -4,3 +4,5 @@ scikit-image
 imageio
 h5py
 tensorboard
+argcomplete
+wandb

res/env/log.json

@@ -1,3 +1,4 @@
 {
-  "summary": false
+  "summary": false,
+  "wandb": false
 }

run.py

@@ -26,16 +26,23 @@ def choices(prefix, parsed_args, **kwargs):
 argcomplete.autocomplete(parser)
 args = parser.parse_args()
 
+
 process = [
     'python3', '-m', 'main',
     '-m', str(args.model_config_path),
     '-d', str(args.dataset_config_path),
     '-r', str(args.run_config_path),
-    '-g', str(args.gpus),
 ]
+gpus = str(args.gpus)
+if gpus == 'cpu':
+    os.environ['CUDA_VISIBLE_DEVICES'] = ''
+    process.extend(['-g', 'cpu'])
+else:
+    os.environ['CUDA_VISIBLE_DEVICES'] = gpus
+    process.extend(['-g', ','.join([str(idx) for idx in range(len(gpus.split(',')))])])
 if args.test_epoch is not None:
     process.append('-t')
     process.append(str(args.test_epoch))
 if args.ci:
     process.append('--ci')
-subprocess.run(process, check=True)
+subprocess.run(process, env=os.environ, check=True)