text: # Copyright (c) 2020 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 paddle.fluid as fluid
  import unittest
  import numpy as np
  import os
  import six
  
  
  def create_reader(shape, batch_number):
      def __impl__():
          idx = 0
          for _ in six.moves.range(batch_number):
              yield np.ones(shape).astype('float32') * idx,
              idx += 1
  
      return __impl__
  
  
  class DataLoaderKeepOrderTestBase(unittest.TestCase):
      def initParameters(self):
          self.iterable = False
          self.break_num = 100
  
      def setUp(self):
          self.epoch_num = 3
          self.batch_num = 40
          self.shape = [3, 4, 5]
          self.initParameters()
  
      def build_network(self, places):
          input_data = fluid.data(shape=self.shape, dtype='float32', name="input")
          loader = fluid.io.DataLoader.from_generator(
              capacity=16, feed_list=[input_data], iterable=self.iterable)
  
          fc = fluid.layers.fc(input_data, size=10)
          loss = fluid.layers.reduce_mean(fc)
  
          loader.set_batch_generator(
              create_reader(self.shape, self.batch_num),
              places=places if loader.iterable else None)
  
          return input_data, loss, loader
  
      def assertInputData(self, batch_id, input_data, dev_cnt):
          if isinstance(input_data, list):
              self.assertTrue(len(input_data), dev_cnt)
              start_val = dev_cnt * batch_id
              for each_input_dict in input_data:
                  input_tensor = np.array(each_input_dict["input"])
                  self.assertEqual(self.shape, list(input_tensor.shape))
                  self.assertTrue((input_tensor == start_val).all())
                  start_val += 1
          else:
              self.assertEqual(
                  list(input_data.shape),
                  [self.shape[0] * dev_cnt] + self.shape[1:])
              start_val = dev_cnt * batch_id
              for idx in six.moves.range(dev_cnt):
                  data_part = input_data[idx * self.shape[0]:(idx + 1) *
                                         self.shape[0], :]
                  self.assertTrue((data_part == start_val).all())
                  start_val += 1
  
      def get_places(self):
          place_list = [fluid.cpu_places(1), fluid.cpu_places(4)]
          if fluid.is_compiled_with_cuda():
              place_list.extend([fluid.cuda_places(0), fluid.cuda_places([0, 1])])
          return place_list
  
      def test_main(self):
          for p in self.get_places():
              use_compiled_program_list = [True] if len(p) > 1 else [False, True]
              for use_compiled_program in use_compiled_program_list:
                  self.run_main_with_place(p, use_compiled_program)
  
      def run_main_with_place(self, places, use_compiled_program=True):
          with fluid.scope_guard(fluid.Scope()):
              with fluid.program_guard(fluid.Program(), fluid.Program()):
                  input_data, loss, loader = self.build_network(places)
                  fetch_list = [input_data]
  
                  exe = fluid.Executor(places[0])
                  exe.run(fluid.default_startup_program())
  
                  dev_cnt = len(places)
                  if dev_cnt > 1:
                      self.assertTrue(use_compiled_program)
  
                  main_program = fluid.default_main_program()
                  if use_compiled_program:
                      main_program = fluid.CompiledProgram(
                          main_program).with_data_parallel(
                              loss_name=loss.name, places=places)
  
                  max_batch_num = min(self.break_num,
                                      int(self.batch_num / dev_cnt))
  
                  if loader.iterable:
                      early_break = False
                      for epoch_id in six.moves.range(self.epoch_num):
                          early_break = False
                          batch_id = 0
                          for data in loader():
                              if batch_id >= self.break_num:
                                  early_break = True
                                  break
                              self.assertInputData(batch_id, data, dev_cnt)
                              fetch_val, = exe.run(program=main_program,
                                                   feed=data,
                                                   fetch_list=fetch_list)
                              self.assertInputData(batch_id, fetch_val, dev_cnt)
                              batch_id += 1
  
                          self.assertEqual(batch_id, max_batch_num)
  
                      if early_break:
                          loader._reset()
                  else:
                      for epoch_id in six.moves.range(self.epoch_num):
                          batch_id = 0
                          loader.start()
                          try:
                              while True:
                                  if batch_id >= self.break_num:
                                      loader.reset()
                                      break
                                  fetch_val, = exe.run(program=main_program,
                                                       fetch_list=fetch_list)
                                  self.assertInputData(batch_id, fetch_val,
                                                       dev_cnt)
                                  batch_id += 1
                          except fluid.core.EOFException:
                              loader.reset()
  
                          self.assertEqual(batch_id, max_batch_num)
  
  
  class IterableDataLoaderKeepOrderTest2(DataLoaderKeepOrderTestBase):
      def initParameters(self):
          self.iterable = True
          self.break_num = 100
  
  
  class IterableDataLoaderKeepOrderTest3(DataLoaderKeepOrderTestBase):
      def initParameters(self):
          self.iterable = False
          self.break_num = 2
  
  
  class IterableDataLoaderKeepOrderTest4(DataLoaderKeepOrderTestBase):
      def initParameters(self):
          self.iterable = True
          self.break_num = 2
  
  
  class IterableDataLoaderKeepOrderTest5(DataLoaderKeepOrderTestBase):
      def initParameters(self):
          self.iterable = False
          self.break_num = 0
  
  
  class IterableDataLoaderKeepOrderTest6(DataLoaderKeepOrderTestBase):
      def initParameters(self):
          self.iterable = True
          self.break_num = 0
  
  
  if __name__ == '__main__':
      unittest.main()
  