test_dataset.py

#   Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import time
import unittest
import sys
import logging
import random
import copy

import set_env

import ppdet.data.transform as tf
from ppdet.data.dataset import Dataset

class MemorySource(Dataset):
    """ memory data source for testing
    """
    def __init__(self, samples):
        super(MemorySource, self).__init__()
        self._epoch = -1

        self._pos = -1
        self._drained = False
        self._samples = samples

    def next(self):
        if self._epoch < 0:
            self.reset()

        if self._pos >= self.size():
            self._drained = True
            raise StopIteration("no more data in " + str(self))
        else:
            sample = copy.deepcopy(self._samples[self._pos])
            self._pos += 1
            return sample

    def reset(self):
        if self._epoch < 0:
            self._epoch = 0
        else:
            self._epoch += 1

        self._pos = 0
        self._drained = False
        random.shuffle(self._samples)

    def size(self):
        return len(self._samples)

    def drained(self):
        assert self._epoch >= 0, "the first epoch has not started yet"
        return self._pos >= self.size()

    def epoch_id(self):
        return self._epoch


class TestDataset(unittest.TestCase):
    """Test cases for ppdet.data.dataset
    """

    @classmethod
    def setUpClass(cls):
        """ setup
        """
        pass

    @classmethod
    def tearDownClass(cls):
        """ tearDownClass """
        pass

    def test_next(self):
        """ test next
        """
        samples = list(range(10))
        mem_sc = MemorySource(samples)

        for i, d in enumerate(mem_sc):
            self.assertTrue(d in samples)

    def test_transform_with_abnormal_worker(self):
        """ test dataset transform with abnormally exit process
        """
        samples = list(range(1000))
        ds = MemorySource(samples)

        def _mapper(sample):
            if sample == 3:
                sys.exit(1)

            return 2 * sample

        worker_conf = {'WORKER_NUM': 2, 'use_process': True}
        mapped = tf.map(ds, _mapper, worker_conf)

        ct = 0
        for i, d in enumerate(mapped):
            ct += 1
            self.assertTrue(d / 2 in samples)

        self.assertEqual(len(samples) - 1, ct)

    def test_transform_with_delay_worker(self):
        """ test dataset transform with delayed process
        """
        samples = list(range(1000))
        ds = MemorySource(samples)

        def _mapper(sample):
            if sample == 3:
                time.sleep(30)

            return 2 * sample

        worker_conf = {'WORKER_NUM': 2, 'use_process': True}
        mapped = tf.map(ds, _mapper, worker_conf)

        ct = 0
        for i, d in enumerate(mapped):
            ct += 1
            self.assertTrue(d / 2 in samples)

        self.assertEqual(len(samples), ct)


if __name__ == '__main__':
    logging.basicConfig()
    unittest.main()