Dev (#240)

* rebase to master

* little modification

* augmentation implemented with torch

* use cuda11.3+torch11, torchrun

* refactor

* support other value of ignore label

* discard distributed

Author: CoinCheung

.gitignore

@@ -114,7 +114,7 @@ tensorrt/build/*
 datasets/coco/train.txt
 datasets/coco/val.txt
 pretrained/*
-dist_train.sh
+run.sh
 openvino/build/*
 openvino/output*
 *.onnx

README.md

@@ -46,11 +46,11 @@ Triton Inference Server(TIS) provides a service solution of deployment. You can
 My platform is like this: 
 
 * ubuntu 18.04
-* nvidia Tesla T4 gpu, driver 450.51.05
-* cuda 10.2
-* cudnn 7
+* nvidia Tesla T4 gpu, driver 450.80.02
+* cuda 10.2/11.3
+* cudnn 8
 * miniconda python 3.8.8
-* pytorch 1.8.1
+* pytorch 1.11.0
 
 
 ## get start
@@ -114,33 +114,7 @@ Then you need to change the field of `im_root` and `train/val_im_anns` in the co
 
 ## train
 
-I used the following command to train the models:
-
-```bash
-# bisenetv1 cityscapes
-export CUDA_VISIBLE_DEVICES=0,1
-cfg_file=configs/bisenetv1_city.py
-NGPUS=2
-python -m torch.distributed.launch --nproc_per_node=$NGPUS tools/train_amp.py --config $cfg_file 
-
-# bisenetv2 cityscapes
-export CUDA_VISIBLE_DEVICES=0,1
-cfg_file=configs/bisenetv2_city.py
-NGPUS=2
-python -m torch.distributed.launch --nproc_per_node=$NGPUS tools/train_amp.py --config $cfg_file 
-
-# bisenetv1 cocostuff
-export CUDA_VISIBLE_DEVICES=0,1,2,3
-cfg_file=configs/bisenetv1_coco.py
-NGPUS=4
-python -m torch.distributed.launch --nproc_per_node=$NGPUS tools/train_amp.py --config $cfg_file 
-
-# bisenetv2 cocostuff
-export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
-cfg_file=configs/bisenetv2_coco.py
-NGPUS=8
-python -m torch.distributed.launch --nproc_per_node=$NGPUS tools/train_amp.py --config $cfg_file 
-```
+Training commands I used to train the models can be found in [here](./dist_train.sh).
 
 Note:  
 1. though `bisenetv2` has fewer flops, it requires much more training iterations. The the training time of `bisenetv1` is shorter.

dist_train.sh

@@ -0,0 +1,32 @@
+
+'''
+NOTE: replace torchrun with torch.distributed.launch if you use older version of pytorch. I suggest you use the same version as I do since I have not tested compatibility with older version after updating.
+'''
+
+
+## bisenetv1 cityscapes
+export CUDA_VISIBLE_DEVICES=0,1
+cfg_file=configs/bisenetv1_city.py
+NGPUS=2
+torchrun --nproc_per_node=$NGPUS tools/train_amp.py --config $cfg_file
+
+
+## bisenetv2 cityscapes
+export CUDA_VISIBLE_DEVICES=0,1
+cfg_file=configs/bisenetv2_city.py
+NGPUS=2
+torchrun --nproc_per_node=$NGPUS tools/train_amp.py --config $cfg_file
+
+
+## bisenetv1 cocostuff
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+cfg_file=configs/bisenetv1_coco.py
+NGPUS=4
+torchrun --nproc_per_node=$NGPUS tools/train_amp.py --config $cfg_file
+
+
+## bisenetv2 cocostuff
+export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+cfg_file=configs/bisenetv2_coco.py
+NGPUS=8
+torchrun --nproc_per_node=$NGPUS tools/train_amp.py --config $cfg_file

lib/base_dataset.py

@@ -24,6 +24,7 @@ def __init__(self, dataroot, annpath, trans_func=None, mode='train'):
         self.mode = mode
         self.trans_func = trans_func
 
+        self.lb_ignore = -100
         self.lb_map = None
 
         with open(annpath, 'r') as fr:
@@ -50,7 +51,8 @@ def __getitem__(self, idx):
         return img.detach(), label.unsqueeze(0).detach()
 
     def get_image(self, impth, lbpth):
-        img, label = cv2.imread(impth)[:, :, ::-1], cv2.imread(lbpth, 0)
+        img = cv2.imread(impth)[:, :, ::-1].copy()
+        label = cv2.imread(lbpth, 0)
         return img, label
 
     def __len__(self):

lib/coco.py

@@ -48,9 +48,9 @@
 class CocoStuff(BaseDataset):
 
     def __init__(self, dataroot, annpath, trans_func=None, mode='train'):
-        super(CocoStuff, self).__init__(dataroot, annpath, trans_func, mode)
+        super(CocoStuff, self).__init__(
+                dataroot, annpath, trans_func, mode)
         self.n_cats = 171 # 91 stuff, 91 thing, 11 of thing have no annos
-        self.lb_ignore = 255
 
         ## label mapping, remove non-existing labels
         missing = [11, 25, 28, 29, 44, 65, 67, 68, 70, 82, 90]

lib/get_dataloader.py

@@ -10,48 +10,25 @@
 
 
 
-class TransformationTrain(object):
 
-    def __init__(self, scales, cropsize):
-        self.trans_func = T.Compose([
-            T.RandomResizedCrop(scales, cropsize),
-            T.RandomHorizontalFlip(),
-            T.ColorJitter(
-                brightness=0.4,
-                contrast=0.4,
-                saturation=0.4
-            ),
-        ])
 
-    def __call__(self, im_lb):
-        im_lb = self.trans_func(im_lb)
-        return im_lb
-
-
-class TransformationVal(object):
-
-    def __call__(self, im_lb):
-        im, lb = im_lb['im'], im_lb['lb']
-        return dict(im=im, lb=lb)
-
-
-def get_data_loader(cfg, mode='train', distributed=True):
+def get_data_loader(cfg, mode='train'):
     if mode == 'train':
-        trans_func = TransformationTrain(cfg.scales, cfg.cropsize)
+        trans_func = T.TransformationTrain(cfg.scales, cfg.cropsize)
         batchsize = cfg.ims_per_gpu
         annpath = cfg.train_im_anns
         shuffle = True
         drop_last = True
     elif mode == 'val':
-        trans_func = TransformationVal()
+        trans_func = T.TransformationVal()
         batchsize = cfg.eval_ims_per_gpu
         annpath = cfg.val_im_anns
         shuffle = False
         drop_last = False
 
     ds = eval(cfg.dataset)(cfg.im_root, annpath, trans_func=trans_func, mode=mode)
 
-    if distributed:
+    if dist.is_initialized():
         assert dist.is_available(), "dist should be initialzed"
         if mode == 'train':
             assert not cfg.max_iter is None

lib/logger.py

@@ -19,6 +19,7 @@ def setup_logger(name, logpth):
     try:
         logging.basicConfig(level=log_level, format=FORMAT, filename=logfile, force=True)
     except Exception:
+        for hl in logging.root.handlers: logging.root.removeHandler(hl)
         logging.basicConfig(level=log_level, format=FORMAT, filename=logfile)
     logging.root.addHandler(logging.StreamHandler())
 

lib/ohem_ce_loss.py

@@ -10,30 +10,30 @@
 #  import ohem_cpp
 #  class OhemCELoss(nn.Module):
 #
-#      def __init__(self, thresh, ignore_lb=255):
+#      def __init__(self, thresh, lb_ignore=255):
 #          super(OhemCELoss, self).__init__()
 #          self.score_thresh = thresh
-#          self.ignore_lb = ignore_lb
-#          self.criteria = nn.CrossEntropyLoss(ignore_index=ignore_lb, reduction='mean')
+#          self.lb_ignore = lb_ignore
+#          self.criteria = nn.CrossEntropyLoss(ignore_index=lb_ignore, reduction='mean')
 #
 #      def forward(self, logits, labels):
-#          n_min = labels[labels != self.ignore_lb].numel() // 16
+#          n_min = labels[labels != self.lb_ignore].numel() // 16
 #          labels = ohem_cpp.score_ohem_label(
-#                  logits, labels, self.ignore_lb, self.score_thresh, n_min).detach()
+#                  logits, labels, self.lb_ignore, self.score_thresh, n_min).detach()
 #          loss = self.criteria(logits, labels)
 #          return loss
 
 
 class OhemCELoss(nn.Module):
 
-    def __init__(self, thresh, ignore_lb=255):
+    def __init__(self, thresh, lb_ignore=255):
         super(OhemCELoss, self).__init__()
         self.thresh = -torch.log(torch.tensor(thresh, requires_grad=False, dtype=torch.float)).cuda()
-        self.ignore_lb = ignore_lb
-        self.criteria = nn.CrossEntropyLoss(ignore_index=ignore_lb, reduction='none')
+        self.lb_ignore = lb_ignore
+        self.criteria = nn.CrossEntropyLoss(ignore_index=lb_ignore, reduction='none')
 
     def forward(self, logits, labels):
-        n_min = labels[labels != self.ignore_lb].numel() // 16
+        n_min = labels[labels != self.lb_ignore].numel() // 16
         loss = self.criteria(logits, labels).view(-1)
         loss_hard = loss[loss > self.thresh]
         if loss_hard.numel() < n_min:

lib/transform_cv2.py

@@ -151,6 +151,30 @@ def __call__(self, im_lb):
         return im_lb
 
 
+class TransformationTrain(object):
+
+    def __init__(self, scales, cropsize):
+        self.trans_func = Compose([
+            RandomResizedCrop(scales, cropsize),
+            RandomHorizontalFlip(),
+            ColorJitter(
+                brightness=0.4,
+                contrast=0.4,
+                saturation=0.4
+            ),
+        ])
+
+    def __call__(self, im_lb):
+        im_lb = self.trans_func(im_lb)
+        return im_lb
+
+
+class TransformationVal(object):
+
+    def __call__(self, im_lb):
+        im, lb = im_lb['im'], im_lb['lb']
+        return dict(im=im, lb=lb)
+
 
 
 if __name__ == '__main__':

openvino/main.cpp

@@ -9,7 +9,7 @@
 #include <inference_engine.hpp>
 
 
-std::string mdpth("../output_v2/model_v2.xml");
+std::string mdpth("../output_v2/model_v2_city.xml");
 std::string device("CPU"); // GNA does not support argmax, my cpu does not has integrated gpu
 std::string impth("../../example.png"); 
 std::string savepth("./res.jpg");