loader report total_batch instead of total_record (tusen-ai#279)

Author: RogerChern

core/detection_input.py

@@ -590,7 +590,7 @@ class Loader(mx.io.DataIter):
 
     def __init__(self, roidb, transform, data_name, label_name, batch_size=1,
                  shuffle=False, num_worker=None, num_collector=None,
-                 worker_queue_depth=None, collector_queue_depth=None, kv=None, valid_count=-1):
+                 worker_queue_depth=None, collector_queue_depth=None, valid_count=-1):
         """
         This Iter will provide roi data to Fast R-CNN network
         :param roidb: must be preprocessed
@@ -600,11 +600,6 @@ def __init__(self, roidb, transform, data_name, label_name, batch_size=1,
         """
         super().__init__(batch_size=batch_size)
 
-        if kv:
-            (self.rank, self.num_worker) = (kv.rank, kv.num_workers)
-        else:
-            (self.rank, self.num_worker) = (0, 1)
-
         # data processing utilities
         if isinstance(transform, dict):
             self.transform = transform["sample"]
@@ -653,8 +648,8 @@ def index(self):
         return self.total_index[:self.valid_count]
 
     @property
-    def total_record(self):
-        return len(self.index) // self.batch_size * self.batch_size
+    def total_batch(self):
+        return len(self.index) // self.batch_size
 
     @property
     def provide_data(self):
@@ -830,8 +825,7 @@ def __init__(self, roidb, transform, data_name, label_name, batch_size=1,
                               num_worker=num_worker,
                               num_collector=num_collector,
                               worker_queue_depth=worker_queue_depth,
-                              collector_queue_depth=collector_queue_depth,
-                              kv=kv)
+                              collector_queue_depth=collector_queue_depth)
             loaders.append(h_loader)
         if len(v_roidb_part) >= batch_size:
             v_loader = Loader(roidb=v_roidb_part,
@@ -844,18 +838,17 @@ def __init__(self, roidb, transform, data_name, label_name, batch_size=1,
                               num_worker=num_worker,
                               num_collector=num_collector,
                               worker_queue_depth=worker_queue_depth,
-                              collector_queue_depth=collector_queue_depth,
-                              kv=kv)
+                              collector_queue_depth=collector_queue_depth)
             loaders.append(v_loader)
         assert len(loaders) > 0, "at least one loader should be constructed"
         self.__loader = SequentialLoader(loaders)
 
     @property
-    def total_record(self):
-        return sum([it.total_record for it in self.__loader.iters])
+    def total_batch(self):
+        return sum([it.total_batch for it in self.__loader.iters])
 
     def __len__(self):
-        return self.total_record
+        return self.total_batch
 
     def __getattr__(self, attr):
         # delegate unknown keys to underlying iterators

detection_infer_speed.py

@@ -38,7 +38,7 @@ def parse_args():
     data_batch = mx.io.DataBatch(data=[data, im_info, im_id, rec_id])
 
     '''
-    there are some conflicts between `mergebn` and `attach_quantized_node` in graph_optimize.py 
+    there are some conflicts between `mergebn` and `attach_quantized_node` in graph_optimize.py
     when mergebn ahead of attach_quantized_node
     such as `Symbol.ComposeKeyword`
     '''
@@ -52,7 +52,7 @@ def parse_args():
         # raise NotImplementedError
         _, out_shape, _ = sym.get_internals().infer_shape(**worker_data_shape)
         out_shape_dictoinary = dict(zip(sym.get_internals().list_outputs(), out_shape))
-        sym = attach_quantize_node(sym, out_shape_dictoinary, pQuant.WeightQuantizeParam, 
+        sym = attach_quantize_node(sym, out_shape_dictoinary, pQuant.WeightQuantizeParam,
                                    pQuant.ActQuantizeParam, pQuant.quantized_op)
     sym.save(pTest.model.prefix + "_infer_speed.json")
 

detection_test.py

@@ -84,10 +84,9 @@ def parse_args():
                         num_worker=4,
                         num_collector=2,
                         worker_queue_depth=2,
-                        collector_queue_depth=2,
-                        kv=None)
+                        collector_queue_depth=2)
 
-        print("total number of images: {}".format(loader.total_record))
+        print("total number of images: {}".format(loader.total_batch))
 
         data_names = [k[0] for k in loader.provide_data]
 
@@ -163,7 +162,7 @@ def data_enqueue(loader, data_queue):
         enqueue_worker.daemon = True
         enqueue_worker.start()
 
-        for _ in range(loader.total_record):
+        for _ in range(loader.total_batch):
             r = result_queue.get()
 
             rid, id, info, cls, box = r

detection_train.py

@@ -186,7 +186,7 @@ def train_net(config):
     eval_metrics = mx.metric.CompositeEvalMetric(metric_list)
 
     # callback
-    batch_end_callback = [callback.Speedometer(train_data.batch_size, frequent=pGen.log_frequency)]
+    batch_end_callback = [callback.Speedometer(train_data.batch_size, len(train_data) * (end_epoch - begin_epoch), frequent=pGen.log_frequency)]
     batch_end_callback += pModel.batch_end_callbacks or []
     epoch_end_callback = callback.do_checkpoint(model_prefix)
     sym.save(model_prefix + ".json")
@@ -196,7 +196,7 @@ def train_net(config):
     base_lr = pOpt.optimizer.lr * kv.num_workers
     lr_factor = pOpt.schedule.lr_factor or 0.1
 
-    iter_per_epoch = len(train_data) // input_batch_size
+    iter_per_epoch = len(train_data)
     total_iter = iter_per_epoch * (end_epoch - begin_epoch)
     lr_iter = [total_iter + it if it < 0 else it for it in lr_iter]
     lr_iter = [it // kv.num_workers for it in lr_iter]

mask_test.py

@@ -86,10 +86,9 @@ def parse_args():
                         num_worker=4,
                         num_collector=2,
                         worker_queue_depth=2,
-                        collector_queue_depth=2,
-                        kv=None)
+                        collector_queue_depth=2)
 
-        print("total number of images: {}".format(loader.total_record))
+        print("total number of images: {}".format(loader.total_batch))
 
         data_names = [k[0] for k in loader.provide_data]
 
@@ -158,7 +157,7 @@ def data_enqueue(loader, data_queue):
         enqueue_worker.daemon = True
         enqueue_worker.start()
 
-        for index in range(loader.total_record):
+        for _ in range(loader.total_batch):
             r = result_queue.get()
 
             rid, id, info, post_cls_score, post_box, post_cls, mask, mask_score = r

rpn_test.py

@@ -83,22 +83,21 @@ def parse_args():
                         num_worker=4,
                         num_collector=2,
                         worker_queue_depth=2,
-                        collector_queue_depth=2,
-                        kv=None)
+                        collector_queue_depth=2)
 
-        print("total number of images: {}".format(loader.total_record))
+        print("total number of images: {}".format(loader.total_batch))
 
         data_names = [k[0] for k in loader.provide_data]
 
         if index_split == 0:
             arg_params, aux_params = load_checkpoint(pTest.model.prefix, pTest.model.epoch)
             if pModel.process_weight is not None:
                 pModel.process_weight(sym, arg_params, aux_params)
-             
-            # merge batch normalization 
+
+            # merge batch normalization
             from utils.graph_optimize import merge_bn
             sym, arg_params, aux_params = merge_bn(sym, arg_params, aux_params)
-            
+
             for i in pKv.gpus:
                 ctx = mx.gpu(i)
                 mod = DetModule(sym, data_names=data_names, context=ctx)
@@ -131,7 +130,7 @@ def data_enqueue(loader, data_queue):
         enqueue_worker.daemon = True
         enqueue_worker.start()
 
-        for _ in range(loader.total_record):
+        for _ in range(loader.total_batch):
             r = result_queue.get()
 
             rid, id, info, box, score = r

utils/callback.py

@@ -4,8 +4,9 @@
 
 
 class Speedometer(object):
-    def __init__(self, batch_size, frequent=50):
+    def __init__(self, batch_size, total_iter, frequent=50):
         self.batch_size = batch_size
+        self.total_iter = total_iter
         self.frequent = frequent
         self.init = False
         self.tic = 0
@@ -23,8 +24,8 @@ def __call__(self, param):
                 speed = self.frequent * self.batch_size / (time.time() - self.tic)
                 if param.eval_metric is not None:
                     name, value = param.eval_metric.get()
-                    s = "Epoch[%d] Batch [%d]\tIter: %d\tLr: %.5f\tSpeed: %.2f samples/sec\tTrain-" % \
-                        (param.epoch, count, param.iter, param.lr, speed)
+                    s = "Epoch[%d] Batch [%d]\tIter: %d/%d\tLr: %.5f\tSpeed: %.2f samples/sec\tTrain-" % \
+                        (param.epoch, count, param.iter, self.total_iter, param.lr, speed)
                     for n, v in zip(name, value):
                         s += "%s=%f,\t" % (n, v)
                     logging.info(s)