diff --git a/lmnet/executor/convert_weight_from_darknet.py b/lmnet/executor/convert_weight_from_darknet.py
index afba40b36..7fa095e97 100644
--- a/lmnet/executor/convert_weight_from_darknet.py
+++ b/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()
 
diff --git a/lmnet/executor/evaluate.py b/lmnet/executor/evaluate.py
index cae5a385c..a7d573b4a 100644
--- a/lmnet/executor/evaluate.py
+++ b/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:
diff --git a/lmnet/executor/export.py b/lmnet/executor/export.py
index 2a9a828ac..aab1fb131 100644
--- a/lmnet/executor/export.py
+++ b/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 = []
diff --git a/lmnet/executor/measure_latency.py b/lmnet/executor/measure_latency.py
index cb2b326e8..2200521f3 100644
--- a/lmnet/executor/measure_latency.py
+++ b/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()
diff --git a/lmnet/executor/predict.py b/lmnet/executor/predict.py
index 5ca95e607..30b59d312 100644
--- a/lmnet/executor/predict.py
+++ b/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"]
diff --git a/lmnet/executor/profile_model.py b/lmnet/executor/profile_model.py
index 8e3e37d3e..0c631074f 100644
--- a/lmnet/executor/profile_model.py
+++ b/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)
diff --git a/lmnet/executor/train.py b/lmnet/executor/train.py
index 1db9286bb..4f6756aee 100644
--- a/lmnet/executor/train.py
+++ b/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)
diff --git a/lmnet/executor/tune_ray.py b/lmnet/executor/tune_ray.py
index bf107b301..8d7320d6c 100644
--- a/lmnet/executor/tune_ray.py
+++ b/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)
 
diff --git a/lmnet/lmnet/datasets/bdd100k.py b/lmnet/lmnet/datasets/bdd100k.py
index f6a061f1d..87abc7b6e 100644
--- a/lmnet/lmnet/datasets/bdd100k.py
+++ b/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
diff --git a/lmnet/lmnet/datasets/camvid.py b/lmnet/lmnet/datasets/camvid.py
index 871c24b50..c0bfce9d5 100644
--- a/lmnet/lmnet/datasets/camvid.py
+++ b/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
diff --git a/lmnet/lmnet/datasets/cifar10.py b/lmnet/lmnet/datasets/cifar10.py
index ec945e3f3..5c4f44756 100644
--- a/lmnet/lmnet/datasets/cifar10.py
+++ b/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
diff --git a/lmnet/lmnet/datasets/cifar100.py b/lmnet/lmnet/datasets/cifar100.py
index 0cba18386..53a7774c5 100644
--- a/lmnet/lmnet/datasets/cifar100.py
+++ b/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
diff --git a/lmnet/lmnet/datasets/cityscapes.py b/lmnet/lmnet/datasets/cityscapes.py
index 44bf7c91d..3524e22ec 100644
--- a/lmnet/lmnet/datasets/cityscapes.py
+++ b/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])
diff --git a/lmnet/lmnet/datasets/dataset_iterator.py b/lmnet/lmnet/datasets/dataset_iterator.py
index 1ddd1cb1e..18dfa3b4f 100644
--- a/lmnet/lmnet/datasets/dataset_iterator.py
+++ b/lmnet/lmnet/datasets/dataset_iterator.py
@@ -22,7 +22,6 @@
 import numpy as np
 import tensorflow as tf
 
-from lmnet.datasets.base import ObjectDetectionBase, SegmentationBase, KeypointDetectionBase
 from lmnet.datasets.tfds import TFDSMixin
 
 _dataset = None
@@ -37,55 +36,36 @@ def _prefetch_setup(dataset, seed, do_shuffle):
         _dataset._shuffle()
 
 
-def _apply_augmentations(dataset, image, label):
+def _apply_augmentations(dataset, sample):
     augmentor = dataset.augmentor
     pre_processor = dataset.pre_processor
 
-    sample = {'image': image}
-
-    if issubclass(dataset.__class__, SegmentationBase):
-        sample['mask'] = label
-    elif issubclass(dataset.__class__, ObjectDetectionBase):
-        sample['gt_boxes'] = label
-    elif issubclass(dataset.__class__, KeypointDetectionBase):
-        sample['joints'] = label
-    else:
-        sample['label'] = label
-
     if callable(augmentor) and dataset.subset == 'train':
         sample = augmentor(**sample)
 
     if callable(pre_processor):
         sample = pre_processor(**sample)
 
-    image = sample['image']
-
-    if issubclass(dataset.__class__, SegmentationBase):
-        label = sample['mask']
-    elif issubclass(dataset.__class__, ObjectDetectionBase):
-        label = sample['gt_boxes']
-    elif issubclass(dataset.__class__, KeypointDetectionBase):
-        label = sample['heatmap']
-    else:
-        label = sample['label']
-
     # FIXME(tokunaga): dataset should not have their own data format
     if dataset.data_format == "NCHW":
-        image = np.transpose(image, [2, 0, 1])
+        sample['image'] = np.transpose(sample['image'], [2, 0, 1])
 
-    return (image, label)
+    return sample
 
 
 def _process_one_data(i):
-    image, label = _dataset[i]
-    return _apply_augmentations(_dataset, image, label)
+    sample = _dataset[i]
+    return _apply_augmentations(_dataset, sample)
+
 
+def _concat_data(sample_list):
 
-def _concat_data(data_list):
-    images, labels = zip(*data_list)
-    images = np.array(images)
-    labels = np.array(labels)
-    return (images, labels)
+    samples_dict = {}
+
+    for key in sample_list[0].keys():
+        samples_dict[key] = np.stack([sample[key] for sample in sample_list], axis=0)
+
+    return samples_dict
 
 
 def _xorshift32(r):
@@ -198,9 +178,9 @@ def read(self):
         """Return batch size data."""
         result = []
         for i in self._gen_ids(self.dataset.batch_size):
-            image, label = self.dataset[i]
-            image, label = _apply_augmentations(self.dataset, image, label)
-            result.append((image, label))
+            sample = self.dataset[i]
+            sample = _apply_augmentations(self.dataset, sample)
+            result.append(sample)
         return _concat_data(result)
 
 
@@ -221,11 +201,12 @@ def __init__(self, dataset):
 
     def read(self):
         """Return batch size data."""
-        result = []
         batch = self.session.run(self.next_batch)
-        for image, label in zip(batch['image'], batch['label']):
-            image, label = _apply_augmentations(self.dataset, image, label)
-            result.append((image, label))
+        result = []
+        for i in range(batch["image"].shape[0]):
+            sample = {key: batch[key][i] for key in batch.keys()}
+            sample = _apply_augmentations(self.dataset, sample)
+            result.append(sample)
         return _concat_data(result)
 
 
@@ -280,10 +261,10 @@ def __iter__(self):
 
     def __next__(self):
         if self.enable_prefetch:
-            (images, labels) = self.prefetch_result_queue.get()
+            samples_dict = self.prefetch_result_queue.get()
         else:
-            images, labels = self.reader.read()
-        return images, labels
+            samples_dict = self.reader.read()
+        return samples_dict
 
     def feed(self):
         return self.__next__()
diff --git a/lmnet/lmnet/datasets/delta_mark.py b/lmnet/lmnet/datasets/delta_mark.py
index a20bd6011..7566042fe 100644
--- a/lmnet/lmnet/datasets/delta_mark.py
+++ b/lmnet/lmnet/datasets/delta_mark.py
@@ -206,7 +206,9 @@ def __getitem__(self, i, type=None):
         label = data_processor.binarize(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
@@ -302,7 +304,10 @@ def __getitem__(self, i, type=None):
         gt_boxes = annotations[i]
         gt_boxes = np.array(gt_boxes)
         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
diff --git a/lmnet/lmnet/datasets/div2k.py b/lmnet/lmnet/datasets/div2k.py
index 4aa0e4b86..3e88c389e 100644
--- a/lmnet/lmnet/datasets/div2k.py
+++ b/lmnet/lmnet/datasets/div2k.py
@@ -45,7 +45,9 @@ def __getitem__(self, i, type=None):
         target_file = self.files[i]
         image = load_image(target_file)
 
-        return image, None
+        sample = {"image": image, "label": None}
+
+        return sample
 
     def __len__(self):
         return self.num_per_epoch
diff --git a/lmnet/lmnet/datasets/ilsvrc_2012.py b/lmnet/lmnet/datasets/ilsvrc_2012.py
index 6eb5ce476..4f51617b3 100644
--- a/lmnet/lmnet/datasets/ilsvrc_2012.py
+++ b/lmnet/lmnet/datasets/ilsvrc_2012.py
@@ -92,7 +92,10 @@ def __getitem__(self, i, type=None):
 
         label = data_processor.binarize(self.annotations[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
diff --git a/lmnet/lmnet/datasets/image_folder.py b/lmnet/lmnet/datasets/image_folder.py
index 6240f38b7..89f9711f0 100644
--- a/lmnet/lmnet/datasets/image_folder.py
+++ b/lmnet/lmnet/datasets/image_folder.py
@@ -115,7 +115,10 @@ def __getitem__(self, i, type=None):
 
         label = data_processor.binarize(label, 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
diff --git a/lmnet/lmnet/datasets/mscoco.py b/lmnet/lmnet/datasets/mscoco.py
index fe6ee7f47..85e261ec1 100644
--- a/lmnet/lmnet/datasets/mscoco.py
+++ b/lmnet/lmnet/datasets/mscoco.py
@@ -123,7 +123,9 @@ def __getitem__(self, i, type=None):
 
         label = self._label_from_image_id(image_id)
 
-        return (image, label)
+        sample = {"image": image, "mask": label}
+
+        return sample
 
     def __len__(self):
         return self.num_per_epoch
@@ -261,7 +263,9 @@ def __getitem__(self, i, type=None):
         gt_boxes = np.array(gt_boxes)
         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
diff --git a/lmnet/lmnet/datasets/mscoco_2017.py b/lmnet/lmnet/datasets/mscoco_2017.py
index 1b8abba1b..a27b84d1b 100644
--- a/lmnet/lmnet/datasets/mscoco_2017.py
+++ b/lmnet/lmnet/datasets/mscoco_2017.py
@@ -157,7 +157,9 @@ def __getitem__(self, item):
 
         cropped_image, joints = self.crop_from_full_image(full_image, box, joints)
 
-        return cropped_image, joints
+        sample = {'image': cropped_image, 'joints': joints}
+
+        return sample
 
     def __len__(self):
         return len(self.files)
diff --git a/lmnet/lmnet/datasets/open_images_v4.py b/lmnet/lmnet/datasets/open_images_v4.py
index c9f363bea..ecec00cd9 100644
--- a/lmnet/lmnet/datasets/open_images_v4.py
+++ b/lmnet/lmnet/datasets/open_images_v4.py
@@ -214,7 +214,9 @@ def __getitem__(self, i, type=None):
 
         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
diff --git a/lmnet/lmnet/datasets/pascalvoc_2007_2012.py b/lmnet/lmnet/datasets/pascalvoc_2007_2012.py
index 6a23598af..ae7bcac63 100644
--- a/lmnet/lmnet/datasets/pascalvoc_2007_2012.py
+++ b/lmnet/lmnet/datasets/pascalvoc_2007_2012.py
@@ -125,7 +125,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
diff --git a/lmnet/lmnet/datasets/pascalvoc_base.py b/lmnet/lmnet/datasets/pascalvoc_base.py
index 5cb2025ee..0616a55ef 100644
--- a/lmnet/lmnet/datasets/pascalvoc_base.py
+++ b/lmnet/lmnet/datasets/pascalvoc_base.py
@@ -220,7 +220,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
diff --git a/lmnet/lmnet/datasets/tfds.py b/lmnet/lmnet/datasets/tfds.py
index 930a2bdc7..56a7e0786 100644
--- a/lmnet/lmnet/datasets/tfds.py
+++ b/lmnet/lmnet/datasets/tfds.py
@@ -64,7 +64,7 @@ def _format_object_detection_record(record, image_size, num_max_boxes):
     gt_boxes = tf.concat([gt_boxes, dummy_boxes], axis=0)
     gt_boxes = tf.slice(gt_boxes, [0, 0], [num_max_boxes, 5])
 
-    return {"image": image, "label": gt_boxes}
+    return {"image": image, "gt_boxes": gt_boxes}
 
 
 class TFDSMixin:
diff --git a/lmnet/lmnet/datasets/widerface.py b/lmnet/lmnet/datasets/widerface.py
index 10c9cb23b..934c5f42f 100644
--- a/lmnet/lmnet/datasets/widerface.py
+++ b/lmnet/lmnet/datasets/widerface.py
@@ -148,7 +148,10 @@ 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
diff --git a/lmnet/lmnet/networks/classification/base.py b/lmnet/lmnet/networks/classification/base.py
index cfaadfc85..19d8faa33 100644
--- a/lmnet/lmnet/networks/classification/base.py
+++ b/lmnet/lmnet/networks/classification/base.py
@@ -39,6 +39,7 @@ def __init__(
         )
 
         self.weight_decay_rate = weight_decay_rate
+        self.placeholders_dict = {}
 
     def placeholders(self):
         """Placeholders.
@@ -62,21 +63,23 @@ def placeholders(self):
             shape=(self.batch_size, self.num_classes),
             name="labels_placeholder")
 
-        return images_placeholder, labels_placeholder
+        self.placeholders_dict["image"] = images_placeholder
+        self.placeholders_dict["label"] = labels_placeholder
 
-    def inference(self, images, is_training):
+    def inference(self, is_training):
         """inference.
 
         Args:
-            images: images tensor. shape is (batch_num, height, width, channel)
-            is_training:
+            is_training: bool.
 
         """
+        images = self.placeholders_dict["image"]
         base = self.base(images, is_training)
         softmax = tf.nn.softmax(base)
 
-        self.output = tf.identity(softmax, name="output")
-        return self.output
+        self.output_tensor = tf.identity(softmax, name="output")
+
+        return self.output_tensor
 
     def _weight_decay_loss(self):
         """L2 weight decay (regularization) loss."""
@@ -91,14 +94,10 @@ def _weight_decay_loss(self):
 
         return tf.add_n(losses) * self.weight_decay_rate
 
-    def loss(self, softmax, labels):
-        """loss.
+    def loss(self):
 
-        Args:
-            softmax: softmaxed tensor from base. shape is (batch_num, num_classes)
-            labels: onehot labels tensor. shape is (batch_num, num_classes)
-
-        """
+        softmax = self.output_tensor
+        labels = self.placeholders_dict["label"]
 
         with tf.name_scope("loss"):
             labels = tf.cast(labels, tf.float32)
@@ -153,7 +152,11 @@ def _heatmaps(self, target_feature_map):
 
         return results
 
-    def summary(self, output, labels=None):
+    def summary(self):
+
+        output = self.output_tensor
+        labels = self.placeholders_dict["label"]
+
         super().summary(output, labels)
 
         images = self.images if self.data_format == 'NHWC' else tf.transpose(self.images, perm=[0, 2, 3, 1])
@@ -190,7 +193,7 @@ def _calc_top_k(self, softmax, labels, k):
         )
         return accuracy_topk, accuracy_topk_update
 
-    def metrics(self, softmax, labels):
+    def metrics(self):
         """metrics.
 
         Args:
@@ -198,6 +201,10 @@ def metrics(self, softmax, labels):
             labels: onehot labels tensor. shape is (batch_num, num_classes)
 
         """
+
+        softmax = self.output_tensor
+        labels = self.placeholders_dict["label"]
+
         with tf.name_scope("metrics_calc"):
             labels = tf.cast(labels, tf.float32)
 
diff --git a/lmnet/lmnet/networks/classification/resnet.py b/lmnet/lmnet/networks/classification/resnet.py
index 2e15b014b..e9a7c74c3 100644
--- a/lmnet/lmnet/networks/classification/resnet.py
+++ b/lmnet/lmnet/networks/classification/resnet.py
@@ -159,7 +159,7 @@ def base(self, images, is_training):
 
         return self.fc
 
-    def loss(self, softmax, labels):
+    def loss(self):
         """loss.
 
         Args:
@@ -167,6 +167,10 @@ def loss(self, softmax, labels):
             labels: onehot labels tensor. shape is (batch_num, num_classes)
 
         """
+
+        softmax = self.output_tensor
+        labels = self.placeholders_dict["label"]
+
         labels = tf.cast(labels, tf.float32)
 
         if self.is_debug:
diff --git a/lmnet/lmnet/networks/object_detection/yolo_v1.py b/lmnet/lmnet/networks/object_detection/yolo_v1.py
index 7e7f9fb27..f00cdf8a5 100644
--- a/lmnet/lmnet/networks/object_detection/yolo_v1.py
+++ b/lmnet/lmnet/networks/object_detection/yolo_v1.py
@@ -51,6 +51,7 @@ def __init__(
         self.boxes_per_cell = boxes_per_cell
         self.leaky_relu_scale = leaky_relu_scale
         self.num_max_boxes = num_max_boxes
+        self.placeholders_dict = {}
 
         self.loss_function = YoloV1Loss(
             is_debug=self.is_debug,
@@ -79,9 +80,13 @@ def placeholders(self):
             shape=(self.batch_size, self.num_max_boxes, 5),
             name="labels_placeholder")
 
-        return images_placeholder, labels_placeholder
+        self.placeholders_dict["image"] = images_placeholder
+        self.placeholders_dict["gt_boxes"] = labels_placeholder
+
+    def summary(self):
+
+        output = self.output_tensor
 
-    def summary(self, output, labels=None):
         predict_classes, predict_confidence, predict_boxes = self._predictions(output)
 
         tf.compat.v1.summary.histogram("predict_classes", predict_classes)
@@ -89,7 +94,11 @@ def summary(self, output, labels=None):
 
         self._summary_predict_boxes(predict_classes, predict_confidence, predict_boxes, threshold=0.05)
 
-    def metrics(self, output, labels, thresholds=[0.3, 0.5, 0.7]):
+    def metrics(self, thresholds=[0.3, 0.5, 0.7]):
+
+        output = self.output_tensor
+        labels = self.placeholders_dict["gt_boxes"]
+
         predict_boxes = self.predict_boxes(output)
 
         metrics_ops_dict = {}
@@ -430,7 +439,7 @@ def _summary_predict_boxes(self, predict_classes, predict_confidence, predict_bo
                     self.image_size,
                 )
 
-    def loss(self, output, gt_boxes, *args):
+    def loss(self, *args):
         """Loss.
 
         Args:
@@ -438,6 +447,10 @@ def loss(self, output, gt_boxes, *args):
                 shape is [batch_size, self.cell_size * self.cell_size * (self.num_classes + self.boxes_per_cell * 5)]
             gt_boxes: ground truth boxes 3D tensor. [batch_size, max_num_boxes, 4(x, y, w, h)].
         """
+
+        output = self.output_tensor
+        gt_boxes = self.placeholders_dict["gt_boxes"]
+
         gt_boxes = self.convert_gt_boxes_xywh_to_cxcywh(gt_boxes)
 
         with tf.name_scope("gt_boxes"):
@@ -453,10 +466,13 @@ def loss(self, output, gt_boxes, *args):
 
         return self.loss_function(predict_classes, predict_confidence, predict_boxes, gt_boxes)
 
-    def inference(self, images, is_training):
+    def inference(self, is_training):
+
+        images = self.placeholders_dict["image"]
         base = self.base(images, is_training)
-        self.output = tf.identity(base, name="output")
-        return self.output
+        self.output_tensor = tf.identity(base, name="output")
+
+        return self.output_tensor
 
     def base(self, images, is_training):
         self.images = images
diff --git a/lmnet/lmnet/networks/object_detection/yolo_v2.py b/lmnet/lmnet/networks/object_detection/yolo_v2.py
index 98c59bb93..b366dd125 100644
--- a/lmnet/lmnet/networks/object_detection/yolo_v2.py
+++ b/lmnet/lmnet/networks/object_detection/yolo_v2.py
@@ -92,6 +92,8 @@ def __init__(
         self.is_dynamic_image_size = is_dynamic_image_size
         self.change_base_output = change_base_output
 
+        self.placeholders_dict = {}
+
         # Assert image size can mod `32`.
         # TODO(wakisaka): Be enable to change `32`. it depends on pooling times.
         assert self.image_size[0] % 32 == 0
@@ -168,9 +170,14 @@ def placeholders(self):
             shape=(self.batch_size, self.num_max_boxes, 5),
             name="labels_placeholder")
 
-        return images_placeholder, labels_placeholder
+        self.placeholders_dict["image"] = images_placeholder
+        self.placeholders_dict["gt_boxes"] = labels_placeholder
+
+    def summary(self):
+
+        output = self.output_tensor
+        labels = self.placeholders_dict["gt_boxes"]
 
-    def summary(self, output, labels=None):
         super().summary(output, labels)
 
         with tf.name_scope("post_process"):
@@ -235,7 +242,11 @@ def summary(self, output, labels=None):
                     data_format=self.data_format,
                 )
 
-    def metrics(self, output, labels, thresholds=[0.3, 0.5, 0.7]):
+    def metrics(self, thresholds=[0.3, 0.5, 0.7]):
+
+        output = self.output_tensor
+        labels = self.placeholders_dict["gt_boxes"]
+
         with tf.compat.v1.variable_scope("calc_metrics"):
             detect_boxes = self.post_process(output)
             metrics_ops_dict = {}
@@ -642,7 +653,7 @@ def _nms(self, formatted_output, iou_threshold, max_output_size, per_class):
 
         return results
 
-    def loss(self, output, gt_boxes, global_step):
+    def loss(self, global_step):
         """Loss.
 
         Args:
@@ -650,6 +661,10 @@ def loss(self, output, gt_boxes, global_step):
                 shape is [batch_size, self.num_cell * self.num_cell * (self.num_classes + self.boxes_per_cell * 5)]
             gt_boxes: ground truth boxes 3D tensor. [batch_size, max_num_boxes, 4(x, y, w, h, class_id)].
         """
+
+        output = self.output_tensor
+        gt_boxes = self.placeholders_dict["gt_boxes"]
+
         if self.change_base_output:
             predict_classes, predict_confidence, predict_boxes = self._split_predictions(output)
         else:
@@ -659,11 +674,14 @@ def loss(self, output, gt_boxes, global_step):
         gt_boxes = self.convert_gt_boxes_xywh_to_cxcywh(gt_boxes)
         return self.loss_function(predict_classes, predict_confidence, predict_boxes, gt_boxes, global_step)
 
-    def inference(self, images, is_training):
+    def inference(self, is_training):
+
+        images = self.placeholders_dict["image"]
+
         tf.compat.v1.summary.histogram("images", images)
         base = self.base(images, is_training)
-        self.output = tf.identity(base, name="output")
-        return self.output
+        self.output_tensor = tf.identity(base, name="output")
+        return self.output_tensor
 
     def _reorg(self, name, inputs, stride, data_format, use_space_to_depth=True, darknet_original=False):
         with tf.name_scope(name):
diff --git a/lmnet/lmnet/networks/segmentation/base.py b/lmnet/lmnet/networks/segmentation/base.py
index 11b31d6ec..607a7ddd9 100755
--- a/lmnet/lmnet/networks/segmentation/base.py
+++ b/lmnet/lmnet/networks/segmentation/base.py
@@ -41,6 +41,7 @@ def __init__(
             self.label_colors = get_color_map(self.num_classes)
         else:
             self.label_colors = label_colors
+        self.placeholders_dict = {}
 
     def placeholders(self):
         shape = (self.batch_size, self.image_size[0], self.image_size[1], 3) \
@@ -54,11 +55,18 @@ def placeholders(self):
             shape=(self.batch_size, self.image_size[0], self.image_size[1]),
             name="labels_placeholder")
 
-        return images_placeholder, labels_placeholder
+        self.placeholders_dict["image"] = images_placeholder
+        self.placeholders_dict["mask"] = labels_placeholder
+
+    def inference(self, is_training):
+
+        images = self.placeholders_dict["image"]
 
-    def inference(self, images, is_training):
         base = self.base(images, is_training)
-        return tf.identity(base, name="output")
+
+        self.output_tensor = tf.identity(base, name="output")
+
+        return self.output_tensor
 
     def _color_labels(self, images, name=""):
         with tf.name_scope(name):
@@ -80,7 +88,10 @@ def _summary_labels(self, labels):
 
         labels = self._color_labels(labels, name="labels_color")
 
-    def summary(self, output, labels=None):
+    def summary(self):
+
+        output = self.output_tensor
+
         output_transposed = output if self.data_format == 'NHWC' else tf.transpose(output, perm=[0, 2, 3, 1])
         images = self.images if self.data_format == 'NHWC' else tf.transpose(self.images, perm=[0, 2, 3, 1])
         tf.summary.image("input", images)
@@ -96,7 +107,11 @@ def summary(self, output, labels=None):
 
         return overlap_output_input, reversed_image
 
-    def metrics(self, output, labels):
+    def metrics(self):
+
+        output = self.output_tensor
+        labels = self.placeholders_dict["mask"]
+
         output_transposed = output if self.data_format == 'NHWC' else tf.transpose(output, perm=[0, 2, 3, 1])
         self._summary_labels(labels)
 
@@ -158,13 +173,17 @@ def __init__(
 
         self.weight_decay_rate = weight_decay_rate
 
-    def loss(self, output, labels):
+    def loss(self):
         """Loss
 
         Args:
            output: Tensor of network output. shape is (batch_size, output_height, output_width, num_classes).
            labels: Tensor of grayscale image gt labels. shape is (batch_size, height, width).
         """
+
+        output = self.output_tensor
+        labels = self.placeholders_dict["mask"]
+
         if self.data_format == 'NCHW':
             output = tf.transpose(output, perm=[0, 2, 3, 1])
         with tf.name_scope("loss"):
diff --git a/lmnet/lmnet/networks/segmentation/lm_bisenet.py b/lmnet/lmnet/networks/segmentation/lm_bisenet.py
index a420da0e8..708e19f6d 100644
--- a/lmnet/lmnet/networks/segmentation/lm_bisenet.py
+++ b/lmnet/lmnet/networks/segmentation/lm_bisenet.py
@@ -318,7 +318,11 @@ def _weight_decay_loss(self):
 
         return tf.add_n(losses) * self.weight_decay_rate
 
-    def loss(self, output, labels):
+    def loss(self):
+
+        output = self.output_tensor
+        labels = self.placeholders_dict["mask"]
+
         x = self.post_process(output)
         context_1 = self.post_process(self.context_1)
         context_2 = self.post_process(self.context_2)
@@ -340,13 +344,14 @@ def loss(self, output, labels):
             tf.compat.v1.summary.scalar("loss", loss)
             return loss
 
-    def summary(self, output, labels=None):
-        x = self.post_process(output)
-        return super().summary(x, labels)
+    def summary(self):
+        return super().summary()
 
-    def metrics(self, output, labels):
-        x = self.post_process(output)
-        return super().metrics(x, labels)
+    def inference(self, is_training):
+
+        self.output_tensor = self.post_process(super().inference(is_training))
+
+        return self.output_tensor
 
     def post_process(self, output):
         with tf.name_scope("post_process"):
diff --git a/lmnet/lmnet/utils/tfds_builders/classification.py b/lmnet/lmnet/utils/tfds_builders/classification.py
index aa05a43fb..51d4e8951 100644
--- a/lmnet/lmnet/utils/tfds_builders/classification.py
+++ b/lmnet/lmnet/utils/tfds_builders/classification.py
@@ -65,7 +65,10 @@ def _split_generators(self, dl_manager):
         return splits
 
     def _generate_examples(self, dataset):
-        for image, label in dataset:
+        for sample_dict in dataset:
+            image = sample_dict["image"]
+            label = sample_dict["label"]
+
             yield {
                 "image": image,
                 "label": label.tolist().index(1)
@@ -74,5 +77,5 @@ def _generate_examples(self, dataset):
     def _num_shards(self, dataset):
         """Decide a number of shards so as not the size of each shard exceeds 256MiB"""
         max_shard_size = 256 * 1024 * 1024  # 256MiB
-        total_size = sum(image.nbytes for image, _ in dataset)
+        total_size = sum(sample_dict["image"].nbytes for sample_dict in dataset)
         return (total_size + max_shard_size - 1) // max_shard_size
diff --git a/lmnet/lmnet/utils/tfds_builders/object_detection.py b/lmnet/lmnet/utils/tfds_builders/object_detection.py
index 647130211..18e2d0603 100644
--- a/lmnet/lmnet/utils/tfds_builders/object_detection.py
+++ b/lmnet/lmnet/utils/tfds_builders/object_detection.py
@@ -68,7 +68,10 @@ def _split_generators(self, dl_manager):
         return splits
 
     def _generate_examples(self, dataset):
-        for image, annotations in dataset:
+        for sample_dict in dataset:
+            image = sample_dict["image"]
+            annotations = sample_dict["gt_boxes"]
+
             height, width, _ = image.shape
 
             objects = [
@@ -93,5 +96,5 @@ def _generate_examples(self, dataset):
     def _num_shards(self, dataset):
         """Decide a number of shards so as not the size of each shard exceeds 256MiB"""
         max_shard_size = 256 * 1024 * 1024  # 256MiB
-        total_size = sum(image.nbytes for image, _ in dataset)
+        total_size = sum(samples_dict["image"].nbytes for samples_dict in dataset)
         return (total_size + max_shard_size - 1) // max_shard_size
diff --git a/lmnet/tests/executor_tests/test_build_tfds.py b/lmnet/tests/executor_tests/test_build_tfds.py
index 5161a65ae..ec5656246 100644
--- a/lmnet/tests/executor_tests/test_build_tfds.py
+++ b/lmnet/tests/executor_tests/test_build_tfds.py
@@ -68,7 +68,10 @@ def test_build_tfds_classification():
     assert validation_dataset.num_per_epoch == validation_data_num
 
     for _ in range(train_data_num):
-        images, labels = train_dataset.feed()
+        samples_dict = train_dataset.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["label"]
 
         assert isinstance(images, np.ndarray)
         assert images.shape[0] == config.BATCH_SIZE
@@ -81,7 +84,10 @@ def test_build_tfds_classification():
         assert labels.shape[1] == train_dataset.num_classes
 
     for _ in range(validation_data_num):
-        images, labels = validation_dataset.feed()
+        samples_dict = validation_dataset.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["label"]
 
         assert isinstance(images, np.ndarray)
         assert images.shape[0] == config.BATCH_SIZE
@@ -129,7 +135,10 @@ def test_build_tfds_object_detection():
     num_max_boxes = train_dataset.num_max_boxes
 
     for _ in range(train_data_num):
-        images, labels = train_dataset.feed()
+        samples_dict = train_dataset.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["label"]
 
         assert isinstance(images, np.ndarray)
         assert images.shape[0] == config.BATCH_SIZE
@@ -143,7 +152,10 @@ def test_build_tfds_object_detection():
         assert labels.shape[2] == 5
 
     for _ in range(validation_data_num):
-        images, labels = validation_dataset.feed()
+        samples_dict = validation_dataset.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["label"]
 
         assert isinstance(images, np.ndarray)
         assert images.shape[0] == config.BATCH_SIZE
diff --git a/lmnet/tests/lmnet_tests/datasets_tests/test_camvid.py b/lmnet/tests/lmnet_tests/datasets_tests/test_camvid.py
index 34ade7228..6a6dc2386 100644
--- a/lmnet/tests/lmnet_tests/datasets_tests/test_camvid.py
+++ b/lmnet/tests/lmnet_tests/datasets_tests/test_camvid.py
@@ -78,25 +78,25 @@ def _show_image_with_annotation(image, label, colors):
 
 def _test_camvid_basic(train_dataset, test_dataset):
     # test train dataset
-    train_image_files, train_label_files = train_dataset.feed()
-    assert train_image_files.shape[0] == 1
-    assert train_label_files.shape[0] == 1
-
-    train_images, train_labels = train_dataset.feed()
-    assert isinstance(train_images, np.ndarray)
-    assert train_images.shape == (1, 360, 480, 3)
-    assert train_labels.shape == (1, 360, 480)
+    samples_dict = train_dataset.feed()
+    assert samples_dict["image"].shape[0] == 1
+    assert samples_dict["mask"].shape[0] == 1
+
+    samples_dict = train_dataset.feed()
+    assert isinstance(samples_dict["image"], np.ndarray)
+    assert samples_dict["image"].shape == (1, 360, 480, 3)
+    assert samples_dict["mask"].shape == (1, 360, 480)
     # _show_image_with_annotation(train_images[0], train_labels[0], colors)
 
     # test test dataset
-    test_image_files, test_label_files = test_dataset.feed()
-    assert test_image_files.shape[0] == 1
-    assert test_label_files.shape[0] == 1
-
-    test_images, test_labels = test_dataset.feed()
-    assert isinstance(test_images, np.ndarray)
-    assert test_images.shape == (1, 360, 480, 3)
-    assert test_labels.shape == (1, 360, 480)
+    samples_dict = test_dataset.feed()
+    assert samples_dict["image"].shape[0] == 1
+    assert samples_dict["mask"].shape[0] == 1
+
+    samples_dict = test_dataset.feed()
+    assert isinstance(samples_dict["image"], np.ndarray)
+    assert samples_dict["image"].shape == (1, 360, 480, 3)
+    assert samples_dict["mask"].shape == (1, 360, 480)
     # _show_image_with_annotation(test_images[0], test_labels[0], colors)
 
 
@@ -112,9 +112,9 @@ def test_camvid():
     colors = train_dataset.label_colors
     assert len(colors) == 12
 
-    train_image_files, train_label_files = train_dataset.feed()
-    assert train_image_files.shape[0] == 1
-    assert train_label_files.shape[0] == 1
+    samples_dict = train_dataset.feed()
+    assert samples_dict["image"].shape[0] == 1
+    assert samples_dict["mask"].shape[0] == 1
 
     _test_camvid_basic(train_dataset, test_dataset)
 
diff --git a/lmnet/tests/lmnet_tests/datasets_tests/test_cityscapes.py b/lmnet/tests/lmnet_tests/datasets_tests/test_cityscapes.py
index caf4c4138..240796a14 100644
--- a/lmnet/tests/lmnet_tests/datasets_tests/test_cityscapes.py
+++ b/lmnet/tests/lmnet_tests/datasets_tests/test_cityscapes.py
@@ -34,6 +34,6 @@ def test_cityscapes():
     colors = train_dataset.label_colors
     assert len(colors) == 34
 
-    train_image_files, train_label_files = train_dataset.feed()
-    assert train_image_files.shape[0] == batch_size
-    assert train_label_files.shape[0] == batch_size
+    samples_dict = train_dataset.feed()
+    assert samples_dict["image"].shape[0] == batch_size
+    assert samples_dict["mask"].shape[0] == batch_size
diff --git a/lmnet/tests/lmnet_tests/datasets_tests/test_dataset_iterator.py b/lmnet/tests/lmnet_tests/datasets_tests/test_dataset_iterator.py
index 06c25e53e..f83c9706e 100644
--- a/lmnet/tests/lmnet_tests/datasets_tests/test_dataset_iterator.py
+++ b/lmnet/tests/lmnet_tests/datasets_tests/test_dataset_iterator.py
@@ -33,18 +33,18 @@ def test_dataset_iterator_batch_size():
     dataset_iterator = DatasetIterator(dataset)
 
     for i in range(0, 10):
-        images, labels = next(dataset_iterator)
-        assert images.shape[0] == batch_size
-        assert labels.shape[0] == batch_size
+        samples_dict = next(dataset_iterator)
+        assert samples_dict["image"].shape[0] == batch_size
+        assert samples_dict["label"].shape[0] == batch_size
 
     batch_size = 32
     dataset = Dummy(subset="train", batch_size=batch_size)
     dataset_iterator = DatasetIterator(dataset)
 
     for i in range(0, 10):
-        images, labels = next(dataset_iterator)
-        assert images.shape[0] == batch_size
-        assert labels.shape[0] == batch_size
+        samples_dict = next(dataset_iterator)
+        assert samples_dict["image"].shape[0] == batch_size
+        assert samples_dict["label"].shape[0] == batch_size
 
 
 def test_dataset_iterator_batch_order():
@@ -56,11 +56,11 @@ def test_dataset_iterator_batch_order():
     prefetch_dataset_iterator = DatasetIterator(dataset, seed=10, enable_prefetch=True)
 
     for i in range(0, 30):
-        images, labels = next(dataset_iterator)
-        prefetch_images, prefetch_labels = next(prefetch_dataset_iterator)
+        samples_dict = next(dataset_iterator)
+        samples_dict_prefetched = next(prefetch_dataset_iterator)
 
-        assert np.all(images == prefetch_images)
-        assert np.all(labels == prefetch_labels)
+        assert np.all(samples_dict["image"] == samples_dict_prefetched["image"])
+        assert np.all(samples_dict["label"] == samples_dict_prefetched["label"])
 
 
 if __name__ == '__main__':
diff --git a/lmnet/tests/lmnet_tests/datasets_tests/test_delta_mark.py b/lmnet/tests/lmnet_tests/datasets_tests/test_delta_mark.py
index 038fc495d..87133e5c4 100644
--- a/lmnet/tests/lmnet_tests/datasets_tests/test_delta_mark.py
+++ b/lmnet/tests/lmnet_tests/datasets_tests/test_delta_mark.py
@@ -84,7 +84,11 @@ def test_delta_mark_classification():
     assert validation_dataset.num_per_epoch == validation_size * all_data_num
 
     for _ in range(5):
-        images, labels = train_dataset.feed()
+
+        samples_dict = train_dataset.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["label"]
 
         assert isinstance(images, np.ndarray)
         assert images.shape[0] == batch_size
@@ -122,7 +126,10 @@ def test_delta_mark_classification_has_validation_path():
     assert validation_dataset.num_per_epoch == validation_data_num
 
     for _ in range(5):
-        images, labels = train_dataset.feed()
+        samples_dict = train_dataset.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["label"]
 
         assert isinstance(images, np.ndarray)
         assert images.shape[0] == batch_size
@@ -135,7 +142,10 @@ def test_delta_mark_classification_has_validation_path():
         assert labels.shape[1] == train_dataset.num_classes
 
     for _ in range(5):
-        images, labels = validation_dataset.feed()
+        samples_dict = validation_dataset.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["label"]
 
         assert isinstance(images, np.ndarray)
         assert images.shape[0] == batch_size
@@ -178,7 +188,11 @@ def test_delta_mark_object_detection():
     assert validation_dataset.num_per_epoch == validation_size * all_data_num
 
     for _ in range(2):
-        images, labels = train_dataset.feed()
+        samples_dict = train_dataset.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["gt_boxes"]
+
         _show_images_with_boxes(images, labels)
 
         assert isinstance(images, np.ndarray)
@@ -222,7 +236,10 @@ def test_delta_mark_object_detection_has_validation_path():
     assert validation_dataset.num_per_epoch == validation_data_num
 
     for _ in range(2):
-        images, labels = train_dataset.feed()
+        samples_dict = train_dataset.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["gt_boxes"]
         # _show_images_with_boxes(images, labels)
 
         assert isinstance(images, np.ndarray)
@@ -237,7 +254,10 @@ def test_delta_mark_object_detection_has_validation_path():
         assert labels.shape[2] == 5
 
     for _ in range(2):
-        images, labels = validation_dataset.feed()
+        samples_dict = validation_dataset.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["gt_boxes"]
         # _show_images_with_boxes(images, labels)
 
         assert isinstance(images, np.ndarray)
diff --git a/lmnet/tests/lmnet_tests/datasets_tests/test_div2k.py b/lmnet/tests/lmnet_tests/datasets_tests/test_div2k.py
index 4f96588ef..a2a638a4b 100644
--- a/lmnet/tests/lmnet_tests/datasets_tests/test_div2k.py
+++ b/lmnet/tests/lmnet_tests/datasets_tests/test_div2k.py
@@ -65,7 +65,11 @@ def test_can_iterate(set_test_environment, subset):
     iterator = DatasetIterator(dataset)
 
     for _ in range(len(dataset)):
-        images, labels = iterator.feed()
+
+        samples_dict = iterator.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["label"]
 
         assert isinstance(images, np.ndarray)
         assert images.shape[0] == batch_size
diff --git a/lmnet/tests/lmnet_tests/datasets_tests/test_image_folder.py b/lmnet/tests/lmnet_tests/datasets_tests/test_image_folder.py
index 5f78151a1..05a48078a 100644
--- a/lmnet/tests/lmnet_tests/datasets_tests/test_image_folder.py
+++ b/lmnet/tests/lmnet_tests/datasets_tests/test_image_folder.py
@@ -60,7 +60,10 @@ def test_image_folder():
     assert len(expected_paths) * (validation_size) == validation_dataset.num_per_epoch
 
     for _ in range(5):
-        images, labels = train_dataset.feed()
+        samples_dict = train_dataset.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["label"]
 
         assert isinstance(images, np.ndarray)
         assert images.shape[0] == batch_size
diff --git a/lmnet/tests/lmnet_tests/datasets_tests/test_mscoco_2017.py b/lmnet/tests/lmnet_tests/datasets_tests/test_mscoco_2017.py
index 6393dd1c3..9de7e13ce 100644
--- a/lmnet/tests/lmnet_tests/datasets_tests/test_mscoco_2017.py
+++ b/lmnet/tests/lmnet_tests/datasets_tests/test_mscoco_2017.py
@@ -42,7 +42,10 @@ def test_mscoco_2017_single_pose_estimation():
     dataset = DatasetIterator(dataset)
 
     for _ in range(5):
-        images, labels = dataset.feed()
+        samples_dict = dataset.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["joints"]
 
         assert isinstance(images, np.ndarray)
         assert images.shape[0] == batch_size
@@ -62,7 +65,10 @@ def test_mscoco_2017_single_pose_estimation():
     dataset = DatasetIterator(dataset)
 
     for _ in range(5):
-        images, labels = dataset.feed()
+        samples_dict = dataset.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["joints"]
 
         assert isinstance(images, np.ndarray)
         assert images.shape[0] == batch_size
diff --git a/lmnet/tests/lmnet_tests/datasets_tests/test_open_images_v4.py b/lmnet/tests/lmnet_tests/datasets_tests/test_open_images_v4.py
index 88e0c77d7..4c28c2068 100644
--- a/lmnet/tests/lmnet_tests/datasets_tests/test_open_images_v4.py
+++ b/lmnet/tests/lmnet_tests/datasets_tests/test_open_images_v4.py
@@ -34,7 +34,10 @@ def test_open_images_v4_classification():
     dataset = DatasetIterator(dataset)
 
     for _ in range(5):
-        images, labels = dataset.feed()
+        samples_dict = dataset.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["label"]
 
         assert isinstance(images, np.ndarray)
         assert images.shape[0] == batch_size
@@ -83,7 +86,10 @@ def test_open_images_v4_object_detection():
     assert dataset.num_max_boxes == OpenImagesV4BoundingBox.count_max_boxes()
 
     for _ in range(5):
-        images, labels = dataset.feed()
+        samples_dict = dataset.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["gt_boxes"]
 
         # _show_images_with_boxes(images, labels)
 
@@ -125,7 +131,10 @@ def test_custom_open_images_v4_object_detection():
     assert validation_dataset.num_per_epoch == 10 * (validation_size)
 
     for _ in range(13):
-        images, labels = train_dataset.feed()
+        samples_dict = train_dataset.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["gt_boxes"]
 
         # _show_images_with_boxes(images, labels)
 
@@ -165,7 +174,10 @@ def test_custom_has_validation_open_images_v4_object_detection():
     assert len(validation_dataset.classes) == 44
 
     for _ in range(3):
-        images, labels = train_dataset.feed()
+        samples_dict = train_dataset.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["gt_boxes"]
 
         # _show_images_with_boxes(images, labels)
 
@@ -181,7 +193,10 @@ def test_custom_has_validation_open_images_v4_object_detection():
         assert labels.shape[2] == 5
 
     for _ in range(3):
-        images, labels = validation_dataset.feed()
+        samples_dict = validation_dataset.feed()
+
+        images = samples_dict["image"]
+        labels = samples_dict["gt_boxes"]
 
         # _show_images_with_boxes(images, labels)
 
diff --git a/lmnet/tests/lmnet_tests/test_data_augmentor.py b/lmnet/tests/lmnet_tests/test_data_augmentor.py
index 889c95b15..5619cbccf 100644
--- a/lmnet/tests/lmnet_tests/test_data_augmentor.py
+++ b/lmnet/tests/lmnet_tests/test_data_augmentor.py
@@ -105,7 +105,10 @@ def test_sequence():
     dataset = DatasetIterator(dataset)
 
     for _ in range(5):
-        images, labels = dataset.feed()
+        samples_dict = dataset.feed()
+        images = samples_dict["image"]
+        labels = samples_dict["gt_boxes"]
+
         _show_images_with_boxes(images, labels)
 
 
@@ -119,7 +122,9 @@ def test_blur():
     dataset = DatasetIterator(dataset)
 
     for _ in range(5):
-        images, labels = dataset.feed()
+        samples_dict = dataset.feed()
+        images = samples_dict["image"]
+        labels = samples_dict["gt_boxes"]
         _show_images_with_boxes(images, labels)
 
 
@@ -133,7 +138,9 @@ def test_brightness():
     dataset = DatasetIterator(dataset)
 
     for _ in range(5):
-        images, labels = dataset.feed()
+        samples_dict = dataset.feed()
+        images = samples_dict["image"]
+        labels = samples_dict["gt_boxes"]
         _show_images_with_boxes(images, labels)
 
 
@@ -147,7 +154,9 @@ def test_color():
     dataset = DatasetIterator(dataset)
 
     for _ in range(5):
-        images, labels = dataset.feed()
+        samples_dict = dataset.feed()
+        images = samples_dict["image"]
+        labels = samples_dict["gt_boxes"]
         _show_images_with_boxes(images, labels)
 
 
@@ -161,7 +170,9 @@ def test_contrast():
     dataset = DatasetIterator(dataset)
 
     for _ in range(5):
-        images, labels = dataset.feed()
+        samples_dict = dataset.feed()
+        images = samples_dict["image"]
+        labels = samples_dict["gt_boxes"]
         _show_images_with_boxes(images, labels)
 
 
@@ -189,7 +200,9 @@ def test_filp_left_right():
     dataset = DatasetIterator(dataset)
 
     for _ in range(5):
-        images, labels = dataset.feed()
+        samples_dict = dataset.feed()
+        images = samples_dict["image"]
+        labels = samples_dict["gt_boxes"]
         _show_images_with_boxes(images, labels)
 
 
@@ -203,7 +216,9 @@ def test_filp_top_bottom():
     dataset = DatasetIterator(dataset)
 
     for _ in range(5):
-        images, labels = dataset.feed()
+        samples_dict = dataset.feed()
+        images = samples_dict["image"]
+        labels = samples_dict["gt_boxes"]
         _show_images_with_boxes(images, labels)
 
 
@@ -217,7 +232,9 @@ def test_hue():
     dataset = DatasetIterator(dataset)
 
     for _ in range(5):
-        images, labels = dataset.feed()
+        samples_dict = dataset.feed()
+        images = samples_dict["image"]
+        labels = samples_dict["gt_boxes"]
         _show_images_with_boxes(images, labels)
 
 
@@ -276,7 +293,9 @@ def test_ssd_random_crop():
     dataset = DatasetIterator(dataset)
 
     for _ in range(5):
-        images, labels = dataset.feed()
+        samples_dict = dataset.feed()
+        images = samples_dict["image"]
+        labels = samples_dict["gt_boxes"]
         _show_images_with_boxes(images, labels)
         assert np.all(labels[:, :, 2] <= 512)
         assert np.all(labels[:, :, 3] <= 256)