train

Author: longcw

cfgs/exps/darknet19_exp1.py

@@ -1,4 +1,4 @@
-exp_name = 'darknet19_voc07trainval_exp1'
+exp_name = 'darknet19_voc07trainval_exp3'
 
 pretrained_fname = 'darknet19.weights.npz'
 

test.py

@@ -22,8 +22,8 @@ def preprocess(fname):
 # hyper-parameters
 # ------------
 imdb_name = cfg.imdb_test
-trained_model = cfg.trained_model
-# trained_model = os.path.join(cfg.train_output_dir, 'darknet19_voc07trainval_exp3_158.h5')
+# trained_model = cfg.trained_model
+trained_model = os.path.join(cfg.train_output_dir, 'darknet19_voc07trainval_exp3_118.h5')
 output_dir = cfg.test_output_dir
 
 max_per_image = 300
@@ -91,7 +91,7 @@ def test_net(net, imdb, max_per_image=300, thresh=0.5, vis=False):
             _t['misc'].clear()
 
         if vis:
-            im2show = yolo_utils.draw_detection(ori_im, bboxes, scores, cls_inds, cfg)
+            im2show = yolo_utils.draw_detection(ori_im, bboxes, scores, cls_inds, cfg, thr=0.1)
             if im2show.shape[0] > 1100:
                 im2show = cv2.resize(im2show, (int(1000. * float(im2show.shape[1]) / im2show.shape[0]), 1000))
             cv2.imshow('test', im2show)

train.py

@@ -2,6 +2,7 @@
 import cv2
 import torch
 import numpy as np
+import datetime
 from torch.multiprocessing import Pool
 
 from darknet import Darknet19
@@ -56,10 +57,12 @@
     else:
         exp = cc.open_experiment(cfg.exp_name)
 
+batch_per_epoch = imdb.batch_per_epoch
 train_loss = 0
 bbox_loss, iou_loss, cls_loss = 0., 0., 0.
 cnt = 0
 t = Timer()
+step_cnt = 0
 for step in range(start_epoch * imdb.batch_per_epoch, cfg.max_epoch * imdb.batch_per_epoch):
     t.tic()
     # batch
@@ -84,15 +87,16 @@
     loss.backward()
     optimizer.step()
     cnt += 1
-
+    step_cnt += 1
     duration = t.toc()
     if step % cfg.disp_interval == 0:
         train_loss /= cnt
         bbox_loss /= cnt
         iou_loss /= cnt
         cls_loss /= cnt
-        print('epoch: %d, step: %d, loss: %.3f, bbox_loss: %.3f, iou_loss: %.3f, cls_loss: %.3f (%.2f s/batch)' % (
-            imdb.epoch, step, train_loss, bbox_loss, iou_loss, cls_loss, duration))
+        print('epoch %d[%d/%d], loss: %.3f, bbox_loss: %.3f, iou_loss: %.3f, cls_loss: %.3f (%.2f s/batch, rest:%s)' % (
+            imdb.epoch, step_cnt, batch_per_epoch, train_loss, bbox_loss, iou_loss, cls_loss, duration,
+            str(datetime.timedelta(seconds=int((batch_per_epoch - step_cnt) * duration)))))
 
         if use_tensorboard and step % cfg.log_interval == 0:
             exp.add_scalar_value('loss_train', train_loss, step=step)
@@ -114,6 +118,6 @@
         save_name = os.path.join(cfg.train_output_dir, '{}_{}.h5'.format(cfg.exp_name, imdb.epoch))
         net_utils.save_net(save_name, net)
         print('save model: {}'.format(save_name))
-
+        step_cnt = 0
 
 imdb.close()

utils/yolo.py

@@ -195,14 +195,16 @@ def get_bbox_targets(images, gt_boxes, cls_inds, dontcares, cfg):
     return bbox_targets, cls_targets
 
 
-def draw_detection(im, bboxes, scores, cls_inds, cfg):
+def draw_detection(im, bboxes, scores, cls_inds, cfg, thr=0.3):
     # draw image
     colors = cfg.colors
     labels = cfg.label_names
 
     imgcv = np.copy(im)
     h, w, _ = imgcv.shape
     for i, box in enumerate(bboxes):
+        if scores[i] < thr:
+            continue
         cls_indx = cls_inds[i]
 
         thick = int((h + w) / 300)