refactor and change dataloader method

Author: CoinCheung

configs/bisenetv1_city.py

@@ -6,12 +6,14 @@
     weight_decay=5e-4,
     warmup_iters=1000,
     max_iter=80000,
+    dataset='CityScapes',
     im_root='./datasets/cityscapes',
     train_im_anns='./datasets/cityscapes/train.txt',
     val_im_anns='./datasets/cityscapes/val.txt',
     scales=[0.75, 2.],
     cropsize=[1024, 1024],
     ims_per_gpu=8,
+    eval_ims_per_gpu=2,
     use_fp16=True,
     use_sync_bn=False,
     respth='./res',

configs/bisenetv2_city.py

@@ -7,12 +7,14 @@
     weight_decay=5e-4,
     warmup_iters = 1000,
     max_iter = 150000,
+    dataset='CityScapes',
     im_root='./datasets/cityscapes',
     train_im_anns='./datasets/cityscapes/train.txt',
     val_im_anns='./datasets/cityscapes/val.txt',
     scales=[0.25, 2.],
     cropsize=[512, 1024],
     ims_per_gpu=8,
+    eval_ims_per_gpu=2,
     use_fp16=True,
     use_sync_bn=True,
     respth='./res',

dist_train.sh

@@ -1,6 +1,9 @@
 
-export CUDA_VISIBLE_DEVICES=0,1
-PORT=52335
+export CUDA_VISIBLE_DEVICES=2,3
+PORT=52330
 NGPUS=2
+cfg_file=configs/bisenetv1_city.py
 
-python -m torch.distributed.launch --nproc_per_node=$NGPUS tools/train_amp.py --config configs/bisenetv1_city.py --port $PORT
+python -m torch.distributed.launch --nproc_per_node=$NGPUS tools/train_amp.py --config $cfg_file --port $PORT
+
+# python -m torch.distributed.launch --nproc_per_node=$NGPUS tools/train.py --config $cfg_file --port $PORT

lib/cityscapes_cv2.py

@@ -12,8 +12,7 @@
 import numpy as np
 
 import lib.transform_cv2 as T
-from lib.sampler import RepeatedDistSampler
-from lib.base_dataset import BaseDataset, TransformationTrain, TransformationVal
+from lib.base_dataset import BaseDataset
 
 
 labels_info = [
@@ -74,48 +73,6 @@ def __init__(self, dataroot, annpath, trans_func=None, mode='train'):
         )
 
 
-def get_data_loader(datapth, annpath, ims_per_gpu, scales, cropsize, max_iter=None, mode='train', distributed=True):
-    if mode == 'train':
-        trans_func = TransformationTrain(scales, cropsize)
-        batchsize = ims_per_gpu
-        shuffle = True
-        drop_last = True
-    elif mode == 'val':
-        trans_func = TransformationVal()
-        batchsize = ims_per_gpu
-        shuffle = False
-        drop_last = False
-
-    ds = CityScapes(datapth, annpath, trans_func=trans_func, mode=mode)
-
-    if distributed:
-        assert dist.is_available(), "dist should be initialzed"
-        if mode == 'train':
-            assert not max_iter is None
-            n_train_imgs = ims_per_gpu * dist.get_world_size() * max_iter
-            sampler = RepeatedDistSampler(ds, n_train_imgs, shuffle=shuffle)
-        else:
-            sampler = torch.utils.data.distributed.DistributedSampler(
-                ds, shuffle=shuffle)
-        batchsampler = torch.utils.data.sampler.BatchSampler(
-            sampler, batchsize, drop_last=drop_last
-        )
-        dl = DataLoader(
-            ds,
-            batch_sampler=batchsampler,
-            num_workers=4,
-            pin_memory=True,
-        )
-    else:
-        dl = DataLoader(
-            ds,
-            batch_size=batchsize,
-            shuffle=shuffle,
-            drop_last=drop_last,
-            num_workers=4,
-            pin_memory=True,
-        )
-    return dl
 
 
 

lib/get_dataloader.py

@@ -0,0 +1,54 @@
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+import torch.distributed as dist
+
+from lib.sampler import RepeatedDistSampler
+from lib.base_dataset import TransformationTrain, TransformationVal
+from lib.cityscapes_cv2 import CityScapes
+
+
+def get_data_loader(cfg, mode='train', distributed=True):
+    if mode == 'train':
+        trans_func = 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()
+        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:
+        assert dist.is_available(), "dist should be initialzed"
+        if mode == 'train':
+            assert not cfg.max_iter is None
+            n_train_imgs = cfg.ims_per_gpu * dist.get_world_size() * cfg.max_iter
+            sampler = RepeatedDistSampler(ds, n_train_imgs, shuffle=shuffle)
+        else:
+            sampler = torch.utils.data.distributed.DistributedSampler(
+                ds, shuffle=shuffle)
+        batchsampler = torch.utils.data.sampler.BatchSampler(
+            sampler, batchsize, drop_last=drop_last
+        )
+        dl = DataLoader(
+            ds,
+            batch_sampler=batchsampler,
+            num_workers=4,
+            pin_memory=True,
+        )
+    else:
+        dl = DataLoader(
+            ds,
+            batch_size=batchsize,
+            shuffle=shuffle,
+            drop_last=drop_last,
+            num_workers=4,
+            pin_memory=True,
+        )
+    return dl

tools/evaluate.py

@@ -22,7 +22,7 @@
 from lib.models import model_factory
 from configs import set_cfg_from_file
 from lib.logger import setup_logger
-from lib.cityscapes_cv2 import get_data_loader
+from lib.get_dataloader import get_data_loader
 
 
 
@@ -184,10 +184,9 @@ def __call__(self, net, dl, n_classes):
 
 
 @torch.no_grad()
-def eval_model(net, ims_per_gpu, im_root, im_anns):
+def eval_model(cfg, net):
     is_dist = dist.is_initialized()
-    dl = get_data_loader(im_root, im_anns, ims_per_gpu, None,
-            None, mode='val', distributed=is_dist)
+    dl = get_data_loader(cfg, mode='val', distributed=is_dist)
     net.eval()
 
     heads, mious = [], []
@@ -251,7 +250,7 @@ def evaluate(cfg, weight_pth):
         )
 
     ## evaluator
-    heads, mious = eval_model(net, 2, cfg.im_root, cfg.val_im_anns)
+    heads, mious = eval_model(cfg, net)
     logger.info(tabulate([mious, ], headers=heads, tablefmt='orgtbl'))
 
 

tools/train.py

@@ -19,7 +19,7 @@
 
 from lib.models import model_factory
 from configs import set_cfg_from_file
-from lib.cityscapes_cv2 import get_data_loader
+from lib.get_dataloader import get_data_loader
 from tools.evaluate import eval_model
 from lib.ohem_ce_loss import OhemCELoss
 from lib.lr_scheduler import WarmupPolyLrScheduler
@@ -133,10 +133,7 @@ def train():
     is_dist = dist.is_initialized()
 
     ## dataset
-    dl = get_data_loader(
-            cfg.im_root, cfg.train_im_anns,
-            cfg.ims_per_gpu, cfg.scales, cfg.cropsize,
-            cfg.max_iter, mode='train', distributed=is_dist)
+    dl = get_data_loader(cfg, mode='train', distributed=is_dist)
 
     ## model
     net, criteria_pre, criteria_aux = set_model()
@@ -202,7 +199,7 @@ def train():
 
     logger.info('\nevaluating the final model')
     torch.cuda.empty_cache()
-    heads, mious = eval_model(net, 2, cfg.im_root, cfg.val_im_anns)
+    heads, mious = eval_model(cfg, net)
     logger.info(tabulate([mious, ], headers=heads, tablefmt='orgtbl'))
 
     return

tools/train_amp.py

@@ -20,7 +20,7 @@
 
 from lib.models import model_factory
 from configs import set_cfg_from_file
-from lib.cityscapes_cv2 import get_data_loader
+from lib.get_dataloader import get_data_loader
 from evaluate import eval_model
 from lib.ohem_ce_loss import OhemCELoss
 from lib.lr_scheduler import WarmupPolyLrScheduler
@@ -122,10 +122,7 @@ def train():
     is_dist = dist.is_initialized()
 
     ## dataset
-    dl = get_data_loader(
-            cfg.im_root, cfg.train_im_anns,
-            cfg.ims_per_gpu, cfg.scales, cfg.cropsize,
-            cfg.max_iter, mode='train', distributed=is_dist)
+    dl = get_data_loader(cfg, mode='train', distributed=is_dist)
 
     ## model
     net, criteria_pre, criteria_aux = set_model()
@@ -187,7 +184,7 @@ def train():
 
     logger.info('\nevaluating the final model')
     torch.cuda.empty_cache()
-    heads, mious = eval_model(net, 2, cfg.im_root, cfg.val_im_anns)
+    heads, mious = eval_model(cfg, net)
     logger.info(tabulate([mious, ], headers=heads, tablefmt='orgtbl'))
 
     return