use outer config.py

Author: CoinCheung

README.md

@@ -29,7 +29,7 @@ My platform is like this:
 ## get start
 With a pretrained weight, you can run inference on an single image like this: 
 ```
-$ python tools/demo.py --model bisenetv2 --weight-path /path/to/your/weights.pth --img-path ./example.png
+$ python tools/demo.py --config configs/bisenetv2_city.py --weight-path /path/to/your/weights.pth --img-path ./example.png
 ```
 This would run inference on the image and save the result image to `./res.jpg`.
 
@@ -64,10 +64,10 @@ In order to train the model, you can run command like this:
 $ export CUDA_VISIBLE_DEVICES=0,1
 
 # if you want to train with apex
-$ python -m torch.distributed.launch --nproc_per_node=2 tools/train.py --model bisenetv2 # or bisenetv1
+$ python -m torch.distributed.launch --nproc_per_node=2 tools/train.py --config configs/bisenetv2_city.py # or bisenetv1
 
 # if you want to train with pytorch fp16 feature from torch 1.6
-$ python -m torch.distributed.launch --nproc_per_node=2 tools/train_amp.py --model bisenetv2 # or bisenetv1
+$ python -m torch.distributed.launch --nproc_per_node=2 tools/train_amp.py --config configs/bisenetv2_city.py # or bisenetv1
 ```
 
 Note that though `bisenetv2` has fewer flops, it requires much more training iterations. The the training time of `bisenetv1` is shorter.
@@ -77,17 +77,17 @@ Note that though `bisenetv2` has fewer flops, it requires much more training ite
 You can also load the trained model weights and finetune from it:
 ```
 $ export CUDA_VISIBLE_DEVICES=0,1
-$ python -m torch.distributed.launch --nproc_per_node=2 tools/train.py --finetune-from ./res/model_final.pth --model bisenetv2 # or bisenetv1
+$ python -m torch.distributed.launch --nproc_per_node=2 tools/train.py --finetune-from ./res/model_final.pth --config ./configs/bisenetv2_city.py # or bisenetv1
 
 # same with pytorch fp16 feature
-$ python -m torch.distributed.launch --nproc_per_node=2 tools/train_amp.py --finetune-from ./res/model_final.pth --model bisenetv2 # or bisenetv1
+$ python -m torch.distributed.launch --nproc_per_node=2 tools/train_amp.py --finetune-from ./res/model_final.pth --config ./configs/bisenetv2_city.py # or bisenetv1
 ```
 
 
 ## eval pretrained models
 You can also evaluate a trained model like this: 
 ```
-$ python tools/evaluate.py --model bisenetv1 --weight-path /path/to/your/weight.pth
+$ python tools/evaluate.py --config configs/bisenetv1_city.py --weight-path /path/to/your/weight.pth
 ```
 
 ## Infer with tensorrt

configs/__init__.py

@@ -1,7 +1,6 @@
 
-from .bisenetv1 import cfg as bisenetv1_cfg
-from .bisenetv2 import cfg as bisenetv2_cfg
 
+import importlib
 
 
 class cfg_dict(object):
@@ -10,7 +9,12 @@ def __init__(self, d):
         self.__dict__ = d
 
 
-cfg_factory = dict(
-    bisenetv1=cfg_dict(bisenetv1_cfg),
-    bisenetv2=cfg_dict(bisenetv2_cfg),
-)
+def set_cfg_from_file(cfg_path):
+    spec = importlib.util.spec_from_file_location('cfg_file', cfg_path)
+    cfg_file = importlib.util.module_from_spec(spec)
+    spec_loader = spec.loader.exec_module(cfg_file)
+    cfg = cfg_file.cfg
+    return cfg_dict(cfg)
+
+
+

configs/bisenetv1.py renamed to configs/bisenetv1_city.py

@@ -1,6 +1,6 @@
 
-export CUDA_VISIBLE_DEVICES=6,7
-PORT=52332
+export CUDA_VISIBLE_DEVICES=0,1
+PORT=52335
 NGPUS=2
 
-python -m torch.distributed.launch --nproc_per_node=$NGPUS tools/train_amp.py --model bisenetv2 --port $PORT
+python -m torch.distributed.launch --nproc_per_node=$NGPUS tools/train_amp.py --config configs/bisenetv1_city.py --port $PORT

configs/bisenetv2.py renamed to configs/bisenetv2_city.py

@@ -10,19 +10,19 @@
 
 import lib.transform_cv2 as T
 from lib.models import model_factory
-from configs import cfg_factory
+from configs import set_cfg_from_file
 
 torch.set_grad_enabled(False)
 np.random.seed(123)
 
 
 # args
 parse = argparse.ArgumentParser()
-parse.add_argument('--model', dest='model', type=str, default='bisenetv2',)
+parse.add_argument('--config', dest='config', type=str, default='configs/bisenetv2.py',)
 parse.add_argument('--weight-path', type=str, default='./res/model_final.pth',)
 parse.add_argument('--img-path', dest='img_path', type=str, default='./example.png',)
 args = parse.parse_args()
-cfg = cfg_factory[args.model]
+cfg = set_cfg_from_file(args.config)
 
 
 palette = np.random.randint(0, 256, (256, 3), dtype=np.uint8)

dist_train.sh

@@ -20,7 +20,7 @@
 import torch.distributed as dist
 
 from lib.models import model_factory
-from configs import cfg_factory
+from configs import set_cfg_from_file
 from lib.logger import setup_logger
 from lib.cityscapes_cv2 import get_data_loader
 
@@ -262,13 +262,14 @@ def parse_args():
     parse.add_argument('--weight-path', dest='weight_pth', type=str,
                        default='model_final.pth',)
     parse.add_argument('--port', dest='port', type=int, default=44553,)
-    parse.add_argument('--model', dest='model', type=str, default='bisenetv2',)
+    parse.add_argument('--config', dest='config', type=str,
+            default='configs/bisenetv2.py',)
     return parse.parse_args()
 
 
 def main():
     args = parse_args()
-    cfg = cfg_factory[args.model]
+    cfg = set_cfg_from_file(args.config)
     if not args.local_rank == -1:
         torch.cuda.set_device(args.local_rank)
         dist.init_process_group(backend='nccl',

tools/demo.py

@@ -6,21 +6,22 @@
 import torch
 
 from lib.models import model_factory
-from configs import cfg_factory
+from configs import set_cfg_from_file
 
 torch.set_grad_enabled(False)
 
 
 parse = argparse.ArgumentParser()
-parse.add_argument('--model', dest='model', type=str, default='bisenetv1',)
+parse.add_argument('--config', dest='config', type=str,
+        default='configs/bisenetv2.py',)
 parse.add_argument('--weight-path', dest='weight_pth', type=str,
         default='model_final.pth')
 parse.add_argument('--outpath', dest='out_pth', type=str,
         default='model.onnx')
 args = parse.parse_args()
 
 
-cfg = cfg_factory[args.model]
+cfg = set_cfg_from_file(args.config)
 if cfg.use_sync_bn: cfg.use_sync_bn = False
 
 net = model_factory[cfg.model_type](19, output_aux=False)

tools/evaluate.py

@@ -18,7 +18,7 @@
 from torch.utils.data import DataLoader
 
 from lib.models import model_factory
-from configs import cfg_factory
+from configs import set_cfg_from_file
 from lib.cityscapes_cv2 import get_data_loader
 from tools.evaluate import eval_model
 from lib.ohem_ce_loss import OhemCELoss
@@ -50,13 +50,13 @@ def parse_args():
     parse = argparse.ArgumentParser()
     parse.add_argument('--local_rank', dest='local_rank', type=int, default=-1,)
     parse.add_argument('--port', dest='port', type=int, default=44554,)
-    parse.add_argument('--model', dest='model', type=str, default='bisenetv2',)
+    parse.add_argument('--config', dest='config', type=str,
+            default='configs/bisenetv2.py',)
     parse.add_argument('--finetune-from', type=str, default=None,)
     return parse.parse_args()
 
 args = parse_args()
-cfg = cfg_factory[args.model]
-
+cfg = set_cfg_from_file(args.config)
 
 
 def set_model():

tools/export_onnx.py

@@ -19,7 +19,7 @@
 import torch.cuda.amp as amp
 
 from lib.models import model_factory
-from configs import cfg_factory
+from configs import set_cfg_from_file
 from lib.cityscapes_cv2 import get_data_loader
 from evaluate import eval_model
 from lib.ohem_ce_loss import OhemCELoss
@@ -44,12 +44,13 @@ def parse_args():
     parse = argparse.ArgumentParser()
     parse.add_argument('--local_rank', dest='local_rank', type=int, default=-1,)
     parse.add_argument('--port', dest='port', type=int, default=44554,)
-    parse.add_argument('--model', dest='model', type=str, default='bisenetv2',)
+    parse.add_argument('--config', dest='config', type=str,
+            default='configs/bisenetv2.py',)
     parse.add_argument('--finetune-from', type=str, default=None,)
     return parse.parse_args()
 
 args = parse_args()
-cfg = cfg_factory[args.model]
+cfg = set_cfg_from_file(args.config)