[WIP] Support multi_labels tasks

lmnet/executor/convert_weight_from_darknet.py

@@ -49,9 +49,9 @@ def convert(config, weight_file):
 
         is_training = tf.constant(False, name="is_training")
 
-        images_placeholder, labels_placeholder = model.placeholders()
+        model.placeholders()
 
-        model.inference(images_placeholder, is_training)
+        model.inference(is_training)
 
         init_op = tf.global_variables_initializer()
 

lmnet/executor/evaluate.py

@@ -86,12 +86,12 @@ def evaluate(config, restore_path):
         global_step = tf.Variable(0, name="global_step", trainable=False)
         is_training = tf.constant(False, name="is_training")
 
-        images_placeholder, labels_placeholder = model.placeholders()
+        model.placeholders()
 
-        output = model.inference(images_placeholder, is_training)
+        model.inference(is_training)
 
-        metrics_ops_dict, metrics_update_op = model.metrics(output, labels_placeholder)
-        model.summary(output, labels_placeholder)
+        metrics_ops_dict, metrics_update_op = model.metrics()
+        model.summary()
 
         summary_op = tf.compat.v1.summary.merge_all()
 
@@ -118,11 +118,10 @@ def evaluate(config, restore_path):
     for test_step in range(test_step_size):
         print("test_step", test_step)
 
-        images, labels = validation_dataset.feed()
-        feed_dict = {
-            images_placeholder: images,
-            labels_placeholder: labels,
-        }
+        samples_dict = validation_dataset.feed()
+
+        feed_dict = {model.placeholders_dict[key]: samples_dict[key]
+                     for key in model.placeholders_dict.keys()}
 
         # Summarize at only last step.
         if test_step == test_step_size - 1:

lmnet/executor/export.py

@@ -93,8 +93,8 @@ def _export(config, restore_path, image_path):
 
         is_training = tf.constant(False, name="is_training")
 
-        images_placeholder, _ = model.placeholders()
-        model.inference(images_placeholder, is_training)
+        model.placeholders()
+        model.inference(is_training)
         init_op = tf.global_variables_initializer()
 
         saver = tf.compat.v1.train.Saver(max_to_keep=50)
@@ -120,7 +120,7 @@ def _export(config, restore_path, image_path):
         image = _pre_process(raw_image, config.PRE_PROCESSOR, config.DATA_FORMAT)
         images = np.expand_dims(image, axis=0)
         feed_dict = {
-            images_placeholder: images,
+            model.placeholders_dict["image"]: images,
         }
 
         all_outputs = []

lmnet/executor/measure_latency.py

@@ -52,8 +52,9 @@ def _measure_time(config, restore_path, step_size):
 
         is_training = tf.constant(False, name="is_training")
 
-        images_placeholder, labels_placeholder = model.placeholders()
-        output = model.inference(images_placeholder, is_training)
+        model.placeholders()
+        model.inference(is_training)
+        output = model.output_tensor
 
         init_op = tf.global_variables_initializer()
 
@@ -72,7 +73,7 @@ def _measure_time(config, restore_path, step_size):
     image = _pre_process(raw_image, config.PRE_PROCESSOR, config.DATA_FORMAT)
     images = np.expand_dims(image, axis=0)
     feed_dict = {
-        images_placeholder: images,
+        model.placeholders_dict["image"]: images,
     }
     output_np = sess.run(output, feed_dict=feed_dict)
     if config.POST_PROCESSOR:
@@ -93,7 +94,7 @@ def _measure_time(config, restore_path, step_size):
         image = _pre_process(raw_image, config.PRE_PROCESSOR, config.DATA_FORMAT)
         images = np.expand_dims(image, axis=0)
         feed_dict = {
-            images_placeholder: images,
+            model.placeholders_dict["image"]: images,
         }
 
         start_only_network = time.time()

lmnet/executor/predict.py

@@ -75,8 +75,8 @@ def _run(input_dir, output_dir, config, restore_path, save_images):
 
         is_training = tf.constant(False, name="is_training")
 
-        images_placeholder, _ = model.placeholders()
-        output_op = model.inference(images_placeholder, is_training)
+        model.placeholders()
+        model.inference(is_training)
 
         init_op = tf.global_variables_initializer()
 
@@ -112,8 +112,8 @@ def _run(input_dir, output_dir, config, restore_path, save_images):
         images, raw_images = _get_images(
             image_files, config.DATASET.PRE_PROCESSOR, config.DATA_FORMAT)
 
-        feed_dict = {images_placeholder: images}
-        outputs = sess.run(output_op, feed_dict=feed_dict)
+        feed_dict = {model.placeholders_dict["image"]: images}
+        outputs = sess.run(model.output_tensor, feed_dict=feed_dict)
 
         if config.POST_PROCESSOR:
             outputs = config.POST_PROCESSOR(outputs=outputs)["outputs"]

lmnet/executor/profile_model.py

@@ -46,8 +46,8 @@ def _profile(config, restore_path, bit, unquant_layers):
 
         is_training = tf.constant(False, name="is_training")
 
-        images_placeholder, _ = model.placeholders()
-        model.inference(images_placeholder, is_training)
+        model.placeholders()
+        model.inference(is_training)
 
         init_op = tf.global_variables_initializer()
         saver = tf.compat.v1.train.Saver(max_to_keep=50)

lmnet/executor/train.py

@@ -106,21 +106,23 @@ def start_training(config):
         global_step = tf.Variable(0, name="global_step", trainable=False)
         is_training_placeholder = tf.compat.v1.placeholder(tf.bool, name="is_training_placeholder")
 
-        images_placeholder, labels_placeholder = model.placeholders()
+        model.placeholders()
 
-        output = model.inference(images_placeholder, is_training_placeholder)
+        placeholders_dict = model.placeholders_dict
+
+        model.inference(is_training_placeholder)
         if config.TASK == Tasks.OBJECT_DETECTION:
-            loss = model.loss(output, labels_placeholder, global_step)
+            loss = model.loss(global_step)
         else:
-            loss = model.loss(output, labels_placeholder)
+            loss = model.loss()
         opt = model.optimizer(global_step)
         if use_horovod:
             # add Horovod Distributed Optimizer
             opt = hvd.DistributedOptimizer(opt)
         train_op = model.train(loss, opt, global_step)
-        metrics_ops_dict, metrics_update_op = model.metrics(output, labels_placeholder)
+        metrics_ops_dict, metrics_update_op = model.metrics()
         # TODO(wakisaka): Deal with many networks.
-        model.summary(output, labels_placeholder)
+        model.summary()
 
         summary_op = tf.compat.v1.summary.merge_all()
 
@@ -213,13 +215,12 @@ def start_training(config):
         progbar.update(last_step)
     for step in range(last_step, max_steps):
 
-        images, labels = train_dataset.feed()
+        samples_dict = train_dataset.feed()
+
+        feed_dict = {placeholders_dict[key]: samples_dict[key]
+                     for key in placeholders_dict.keys()}
 
-        feed_dict = {
-            is_training_placeholder: True,
-            images_placeholder: images,
-            labels_placeholder: labels,
-        }
+        feed_dict[is_training_placeholder] = True
 
         if step * ((step + 1) % config.SUMMARISE_STEPS) == 0 and rank == 0:
             # Runtime statistics for develop.
@@ -261,12 +262,12 @@ def start_training(config):
                 for train_validation_saving_step in range(train_validation_saving_step_size):
                     print("train_validation_saving_step", train_validation_saving_step)
 
-                    images, labels = train_validation_saving_dataset.feed()
-                    feed_dict = {
-                        is_training_placeholder: False,
-                        images_placeholder: images,
-                        labels_placeholder: labels,
-                    }
+                    samples_dict = train_dataset.feed()
+
+                    feed_dict = {placeholders_dict[key]: samples_dict[key]
+                                 for key in placeholders_dict.keys()}
+
+                    feed_dict[is_training_placeholder] = False
 
                     if train_validation_saving_step % config.SUMMARISE_STEPS == 0:
                         summary, _ = sess.run([summary_op, metrics_update_op], feed_dict=feed_dict)
@@ -315,12 +316,12 @@ def start_training(config):
 
             for test_step in range(test_step_size):
 
-                images, labels = validation_dataset.feed()
-                feed_dict = {
-                    is_training_placeholder: False,
-                    images_placeholder: images,
-                    labels_placeholder: labels,
-                }
+                samples_dict = train_dataset.feed()
+
+                feed_dict = {placeholders_dict[key]: samples_dict[key]
+                             for key in placeholders_dict.keys()}
+
+                feed_dict[is_training_placeholder] = False
 
                 if test_step % config.SUMMARISE_STEPS == 0:
                     summary, _ = sess.run([summary_op, metrics_update_op], feed_dict=feed_dict)

lmnet/executor/tune_ray.py

@@ -194,17 +194,18 @@ def _setup(self, config):
 
         self.global_step = tf.Variable(0, name="global_step", trainable=False)
         self.is_training_placeholder = tf.compat.v1.placeholder(tf.bool, name="is_training_placeholder")
-        self.images_placeholder, self.labels_placeholder = model.placeholders()
+        model.placeholders()
+        self.placeholders_dict = model.placeholders_dict
 
-        output = model.inference(self.images_placeholder, self.is_training_placeholder)
+        output = model.inference(self.is_training_placeholder)
         if model_class.__module__.startswith("lmnet.networks.object_detection"):
-            loss = model.loss(output, self.labels_placeholder, self.is_training_placeholder)
+            loss = model.loss(self.is_training_placeholder)
         else:
-            loss = model.loss(output, self.labels_placeholder)
+            loss = model.loss()
         opt = model.optimizer(self.global_step)
 
         train_op = model.train(loss, opt, self.global_step)
-        metrics_ops_dict, metrics_update_op = model.metrics(output, self.labels_placeholder)
+        metrics_ops_dict, metrics_update_op = model.metrics(self.labels_placeholder)
 
         self.train_op = train_op
         self.metrics_ops_dict = metrics_ops_dict
@@ -224,25 +225,24 @@ def _train(self):
         step_per_epoch = int(self.train_dataset.num_per_epoch / self.lm_config.BATCH_SIZE)
 
         for _ in range(step_per_epoch):
-            images, labels = self.train_dataset.feed()
+            samples_dict = self.train_dataset.feed()
 
-            feed_dict = {
-                self.is_training_placeholder: True,
-                self.images_placeholder: images,
-                self.labels_placeholder: labels,
-            }
+            feed_dict = {self.placeholders_dict[key]: samples_dict[key]
+                         for key in self.placeholders_dict.keys()}
+
+            feed_dict[self.is_training_placeholder] = True
 
             self.sess.run([self.train_op], feed_dict=feed_dict)
 
         self.sess.run(self.reset_metrics_op)
         test_step_size = int(math.ceil(self.validation_dataset.num_per_epoch / self.lm_config.BATCH_SIZE))
         for _ in range(test_step_size):
-            images, labels = self.validation_dataset.feed()
-            feed_dict = {
-                self.is_training_placeholder: False,
-                self.images_placeholder: images,
-                self.labels_placeholder: labels,
-            }
+            samples_dict = self.validation_dataset.feed()
+
+            feed_dict = {self.placeholders_dict[key]: samples_dict[key]
+                         for key in self.placeholders_dict.keys()}
+
+            feed_dict[self.is_training_placeholder] = False
 
             self.sess.run([self.metrics_update_op], feed_dict=feed_dict)
 

lmnet/lmnet/datasets/bdd100k.py

@@ -149,7 +149,9 @@ def __getitem__(self, i, type=None):
         gt_boxes = gt_boxes.copy()  # is it really needed?
         gt_boxes = self._fill_dummy_boxes(gt_boxes)
 
-        return image, gt_boxes
+        sample = {"image": image, "gt_boxes": gt_boxes}
+
+        return sample
 
     def __len__(self):
         return self.num_per_epoch

lmnet/lmnet/datasets/camvid.py

@@ -106,7 +106,9 @@ def __getitem__(self, i, type=None):
         image = get_image(image_files[i])
         label = get_image(label_files[i], convert_rgb=False, ignore_class_idx=self.ignore_class_idx).copy()
 
-        return (image, label)
+        sample = {"image": image, "mask": label}
+
+        return sample
 
     def __len__(self):
         return self.num_per_epoch

lmnet/lmnet/datasets/cifar10.py

@@ -131,7 +131,10 @@ def __getitem__(self, i, type=None):
         image = self._get_image(i)
         label = data_processor.binarize(self.labels[i], self.num_classes)
         label = np.reshape(label, (self.num_classes))
-        return (image, label)
+
+        sample = {"image": image, "label": label}
+
+        return sample
 
     def __len__(self):
         return self.num_per_epoch

lmnet/lmnet/datasets/cifar100.py

@@ -121,7 +121,10 @@ def __getitem__(self, i, type=None):
         image = self._get_image(i)
         label = data_processor.binarize(self.labels[i], self.num_classes)
         label = np.reshape(label, (self.num_classes))
-        return (image, label)
+
+        sample = {"image": image, "label": label}
+
+        return sample
 
     def __len__(self):
         return self.num_per_epoch

lmnet/lmnet/datasets/cityscapes.py

@@ -135,10 +135,12 @@ def files_and_annotations(self):
 
     def __getitem__(self, i):
         imgs, labels = self.files_and_annotations()
-        img = Image.open(imgs[i])
-        label = Image.open(labels[i])
+        img = np.array(Image.open(imgs[i]))
+        label = np.array(Image.open(labels[i]))
 
-        return np.array(img), np.array(label)
+        sample = {"image": img, "mask": label}
+
+        return sample
 
     def __len__(self):
         return len(self.files_and_annotations()[0])