Update training procedure

kazuto1011 · Nov 3, 2017 · 9741062 · 9741062
1 parent 309cd91
commit 9741062
Show file tree

Hide file tree

Showing 24 changed files with 107 additions and 115 deletions.
diff --git a/config/default.yaml b/config/default.yaml
@@ -1,25 +1,26 @@
 num_workers: 0
-mean:
-    R: 122.675
-    G: 116.669
-    B: 104.008
+image:
+    mean:
+        R: 122.675
+        G: 116.669
+        B: 104.008
+    size:
+        train: 321
+        test: 513
 dataset:
     voc:
+        name: Pascal VOC 2012
         root:
         label_list: /Users/kazuto1011/Desktop/voc-labels.txt
-        init_model: ./MS_DeepLab_resnet_pretrained_COCO_init.pth
+        init_model: ./data/models/MS_DeepLab_resnet_pretrained_COCO_init.pth
         ignore_label: 255
-        rows: 513
-        cols: 513
         n_classes: 21
         trained_model: ./MS_DeepLab_resnet_trained_VOC.pth
     cocostuff:
+        name: COCO Stuff 10k
         root: /Users/kazuto1011/Desktop/cocostuff-10k-v1.1
         label_list: /Users/kazuto1011/Desktop/cocostuff-10k-v1.1/cocostuff-labels.txt
-        init_model: ./MS_DeepLab_resnet_pretrained_COCO_init.pth
+        init_model: ./data/models/MS_DeepLab_resnet_pretrained_COCO_init.pth
         ignore_label: 0
-        rows: 513
-        cols: 513
         n_classes: 183
-        classes:
         trained_model: None
diff --git a/data/results/.gitignore b/data/results/.gitignore
@@ -0,0 +1,2 @@
+*.png
+*.jpg
diff --git a/samples/cat_dog.png → data/samples/cat_dog.png b/samples/cat_dog.png → data/samples/cat_dog.png
diff --git a/samples/coco.jpg → data/samples/coco.jpg b/samples/coco.jpg → data/samples/coco.jpg
diff --git a/samples/coco_2.jpg → data/samples/coco_2.jpg b/samples/coco_2.jpg → data/samples/coco_2.jpg
diff --git a/samples/giraffe.jpg → data/samples/giraffe.jpg b/samples/giraffe.jpg → data/samples/giraffe.jpg
diff --git a/samples/longbeach_grayscale.png → data/samples/longbeach_grayscale.png b/samples/longbeach_grayscale.png → data/samples/longbeach_grayscale.png
diff --git a/samples/longbeach_rgb.jpg → data/samples/longbeach_rgb.jpg b/samples/longbeach_rgb.jpg → data/samples/longbeach_rgb.jpg
diff --git a/demo.py b/demo.py
@@ -6,6 +6,7 @@
 # Created:  2017-11-01
 
 import argparse
+import os.path as osp
 
 import cv2
 import matplotlib.pyplot as plt
@@ -16,15 +17,15 @@
 import yaml
 from torch.autograd import Variable
 
-from models import DeepLab
-import os.path as osp
+from libs.models import DeepLab
 
 
 def main(args):
     # Configuration
     with open(args.config) as f:
         config = yaml.load(f)
 
+    # Label list
     with open(config['dataset'][args.dataset]['label_list']) as f:
         classes = {}
         for label in f:
@@ -49,16 +50,16 @@ def main(args):
     if args.cuda:
         model.cuda()
 
-    image_size = (config['dataset'][args.dataset]['rows'],
-                  config['dataset'][args.dataset]['cols'])
+    image_size = (config['image']['size']['test'],
+                  config['image']['size']['test'])
 
     # Image preprocessing
     image = cv2.imread(args.image, cv2.IMREAD_COLOR).astype(float)
     image = cv2.resize(image, image_size)
     image_original = image.astype(np.uint8)[:, :, ::-1]
-    image -= np.array([config['mean']['B'],
-                       config['mean']['G'],
-                       config['mean']['R']])
+    image -= np.array([config['image']['mean']['B'],
+                       config['image']['mean']['G'],
+                       config['image']['mean']['R']])
     image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0)
     image = image.cuda() if args.cuda else image
 
@@ -90,7 +91,6 @@ def main(args):
         ax.set_xticks([])
         ax.set_yticks([])
 
-    # plt.savefig('./results/{}'.format(osp.basename(args.image)))
     plt.show()
 
 

diff --git a/demo_batch.py b/demo_batch.py
@@ -18,7 +18,7 @@
 import yaml
 from torch.autograd import Variable
 
-from models import DeepLab
+from libs.models import DeepLab
 
 
 def main(args):
@@ -50,8 +50,8 @@ def main(args):
     if args.cuda:
         model.cuda()
 
-    image_size = (config['dataset'][args.dataset]['rows'],
-                  config['dataset'][args.dataset]['cols'])
+    image_size = (config['image']['size']['test'],
+                  config['image']['size']['test'])
 
     image_list = glob(osp.join(args.image_dir, '*'))
 
@@ -60,9 +60,9 @@ def main(args):
         image = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(float)
         image = cv2.resize(image, image_size)
         image_original = image.astype(np.uint8)[:, :, ::-1]
-        image -= np.array([config['mean']['B'],
-                           config['mean']['G'],
-                           config['mean']['R']])
+        image -= np.array([config['image']['mean']['B'],
+                           config['image']['mean']['G'],
+                           config['image']['mean']['R']])
         image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0)
         image = image.cuda() if args.cuda else image
 
@@ -94,8 +94,7 @@ def main(args):
             ax.set_xticks([])
             ax.set_yticks([])
 
-        plt.savefig('./results/{}'.format(osp.basename(image_path)))
-        # plt.show()
+        plt.savefig('./data/results/{}'.format(osp.basename(image_path)))
 
 
 if __name__ == '__main__':

diff --git a/test_.py → eval.py b/test_.py → eval.py
diff --git a/libs/__init__.py b/libs/__init__.py
diff --git a/datasets/__init__.py → libs/datasets/__init__.py b/datasets/__init__.py → libs/datasets/__init__.py
@@ -1,4 +1,4 @@
-from datasets.cocostuff import CocoStuff10k
+from libs.datasets.cocostuff import CocoStuff10k
 
 
 def get_dataset(name):

diff --git a/datasets/cocostuff.py → libs/datasets/cocostuff.py b/datasets/cocostuff.py → libs/datasets/cocostuff.py
diff --git a/libs/models/__init__.py b/libs/models/__init__.py
@@ -0,0 +1 @@
+from libs.models.deeplab import DeepLab
diff --git a/models/deeplab.py → libs/models/deeplab.py b/models/deeplab.py → libs/models/deeplab.py
diff --git a/libs/utils/__init__.py b/libs/utils/__init__.py
@@ -0,0 +1,2 @@
+from libs.utils.loss import CrossEntropyLoss2d
+from libs.utils.lr_scheduler import poly_lr_scheduler
diff --git a/utils/loss.py → libs/utils/loss.py b/utils/loss.py → libs/utils/loss.py
diff --git a/utils/lr_scheduler.py → libs/utils/lr_scheduler.py b/utils/lr_scheduler.py → libs/utils/lr_scheduler.py
diff --git a/models/__init__.py b/models/__init__.py
diff --git a/results/.gitignore b/results/.gitignore
diff --git a/publish_tensorboard.sh → tensorboard.sh b/publish_tensorboard.sh → tensorboard.sh
diff --git a/train.py b/train.py
@@ -17,9 +17,9 @@
 from torchnet.meter import MovingAverageValueMeter
 from tqdm import tqdm
 
-from datasets import get_dataset
-from models import DeepLab
-from utils import CrossEntropyLoss2d
+from libs.datasets import get_dataset
+from libs.models import DeepLab
+from libs.utils import CrossEntropyLoss2d
 
 
 def get_1x_lr_params(model):
@@ -71,9 +71,6 @@ def main(args):
     with open(args.config) as f:
         config = yaml.load(f)
 
-    # TensorBoard Logger
-    writer = SummaryWriter(args.log_dir)
-
     # Dataset
     dataset = get_dataset(args.dataset)(
         root=config['dataset'][args.dataset]['root'],
@@ -91,6 +88,7 @@ def main(args):
         num_workers=config['num_workers'],
         shuffle=True
     )
+    loader_iter = iter(loader)
 
     # Model
     model = DeepLab(n_classes=config['dataset'][args.dataset]['n_classes'])
@@ -120,86 +118,77 @@ def main(args):
 
     # Loss definition
     criterion = CrossEntropyLoss2d(
-        ignore_index=config['dataset'][args.dataset]['ignore_label'])
+        ignore_index=config['dataset'][args.dataset]['ignore_label']
+    )
     if args.cuda:
         criterion.cuda()
 
-    best_loss = 1e10
-    patience = args.patience
+    # TensorBoard Logger
+    writer = SummaryWriter(args.log_dir)
+
+    loss = 0
     loss_meter = MovingAverageValueMeter(20)
 
-    # Training loop
-    iteration = 0
-    for epoch in range(args.n_epoch):
-
-        tqdm_loader = tqdm(
-            enumerate(loader),
-            total=len(loader),
-            desc='Epoch [%d]' % (epoch),
-            leave=False,
-            dynamic_ncols=True
-        )
-
-        model.train()
-        for i, (data, target) in tqdm_loader:
-            # Image
-            data = data.cuda() if args.cuda else data
-            data = Variable(data)
-
-            # Polynomial lr decay
-            poly_lr_scheduler(optimizer=optimizer,
-                              init_lr=args.lr,
-                              iter=iteration,
-                              lr_decay_iter=args.lr_decay,
-                              max_iter=len(loader) * args.n_epoch / args.batch_size,  # NOQA
-                              power=args.poly_power)
-
-            # Forward propagation
+    model.train()
+    for iteration in tqdm(range(1, args.iter_max + 1),
+                          total=args.iter_max,
+                          leave=False,
+                          dynamic_ncols=True):
+
+        data, target = next(loader_iter)
+
+        # Polynomial lr decay
+        poly_lr_scheduler(optimizer=optimizer,
+                          init_lr=args.lr,
+                          iter=iteration - 1,
+                          lr_decay_iter=args.lr_decay,
+                          max_iter=args.iter_max,
+                          power=args.poly_power)
+
+        # Image
+        data = data.cuda() if args.cuda else data
+        data = Variable(data)
+
+        # Forward propagation
+        outputs = model(data)
+
+        # Label
+        target = resize_target(target, outputs[0].size(2))
+        target = target.cuda() if args.cuda else target
+        target = Variable(target)
+
+        # Aggregate losses for [100%, 75%, 50%, Max]
+        for output in outputs:
+            loss += criterion(output, target)
+        loss /= args.iter_size
+        loss_meter.add(loss.data[0])
+
+        # Back propagation
+        if iteration % args.iter_size == 0:
             optimizer.zero_grad()
-            outputs = model(data)
-
-            # Label
-            target = resize_target(target, outputs[0].size(2))
-            target = target.cuda() if args.cuda else target
-            target = Variable(target)
-
-            # Loss
-            loss = 0
-            for output in outputs:
-                loss += criterion(output, target)
-            loss_meter.add(loss.data[0])
-
-            # Back propagation & weight updating
             loss.backward()
             optimizer.step()
+            loss = 0
+
+        # TensorBoard
+        if iteration % args.iter_tf == 0:
+            writer.add_scalar('train_loss', loss_meter.value()[0], iteration)
+
+        # Save a model
+        if iteration % args.iter_snapshot == 0:
+            torch.save(
+                {'iteration': iteration,
+                 'weight': model.state_dict()},
+                osp.join(args.save_dir, 'checkpoint_{}.pth.tar'.format(iteration))
+            )
+            writer.add_text('log', 'Saved a model', iteration)
 
-            # TensorBoard: Scalar
-            train_loss = loss_meter.value()[0]
-            writer.add_scalar('train_loss', train_loss, iteration)
-
-            # In each 1000 iterations
-            if iteration % 100 == 0:
-                # Early stopping procedure
-                if train_loss < best_loss:
-                    torch.save(
-                        {'iteration': iteration,
-                         'weight': model.state_dict()},
-                        osp.join(args.save_dir, 'checkpoint_best.pth.tar')
-                    )
-                    writer.add_text('log', 'Saved a model', iteration)
-                    best_loss = train_loss
-                    patience = args.patience
-                else:
-                    patience -= 1
-                    if patience == 0:
-                        writer.add_text('log', 'Early stopping', iteration)
-                    #     break
-
-            iteration += 1
+        if iteration % len(loader) == 0:
+            loader_iter = iter(loader)
 
     torch.save(
         {'iteration': iteration - 1,
-            'weight': model.state_dict()},
+         'weight': model.state_dict()},
         osp.join(args.save_dir, 'checkpoint_final.pth.tar')
     )
 
@@ -208,16 +197,18 @@ def main(args):
     # Parsing arguments
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--no_cuda', action='store_true', default=False)
-    parser.add_argument('--dataset', nargs='?', type=str, default='cocostuff')
+    parser.add_argument('--dataset', type=str, default='cocostuff')
     parser.add_argument('--config', type=str, default='config/default.yaml')
-    parser.add_argument('--n_epoch', nargs='?', type=int, default=2)
-    parser.add_argument('--batch_size', nargs='?', type=int, default=2)
-    parser.add_argument('--lr', nargs='?', type=float, default=0.00025)
-    parser.add_argument('--lr_decay', nargs='?', type=int, default=10)
-    parser.add_argument('--momentum', nargs='?', type=float, default=0.9)
-    parser.add_argument('--weight_decay', nargs='?', type=float, default=5e-4)
-    parser.add_argument('--poly_power', nargs='?', type=float, default=0.9)
-    parser.add_argument('--patience', type=int, default=10)
+    parser.add_argument('--batch_size', type=int, default=1)
+    parser.add_argument('--lr', type=float, default=2.5e-4)
+    parser.add_argument('--lr_decay', type=int, default=10)
+    parser.add_argument('--momentum', type=float, default=0.9)
+    parser.add_argument('--weight_decay', type=float, default=5e-4)
+    parser.add_argument('--poly_power', type=float, default=0.9)
+    parser.add_argument('--iter_max', type=int, default=20000)
+    parser.add_argument('--iter_size', type=int, default=10)
+    parser.add_argument('--iter_tf', type=int, default=50)
+    parser.add_argument('--iter_snapshot', type=int, default=5000)
     parser.add_argument('--optimizer', type=str, default='sgd')
     parser.add_argument('--save_dir', type=str, default='.')
     parser.add_argument('--log_dir', type=str, default='runs')

diff --git a/utils/__init__.py b/utils/__init__.py
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		from libs.utils.loss import CrossEntropyLoss2d
		from libs.utils.lr_scheduler import poly_lr_scheduler