mix_dataset.py 1.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849
  1. # Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from __future__ import print_function
  15. import numpy as np
  16. import os
  17. from paddle.io import Dataset
  18. from .. import dataloader
  19. class MixDataset(Dataset):
  20. def __init__(self, datasets_config):
  21. super().__init__()
  22. self.dataset_list = []
  23. start_idx = 0
  24. end_idx = 0
  25. for config_i in datasets_config:
  26. dataset_name = config_i.pop('name')
  27. dataset = getattr(dataloader, dataset_name)(**config_i)
  28. end_idx += len(dataset)
  29. self.dataset_list.append([end_idx, start_idx, dataset])
  30. start_idx = end_idx
  31. self.length = end_idx
  32. def __getitem__(self, idx):
  33. for dataset_i in self.dataset_list:
  34. if dataset_i[0] > idx:
  35. dataset_i_idx = idx - dataset_i[1]
  36. return dataset_i[2][dataset_i_idx]
  37. def __len__(self):
  38. return self.length
  39. def get_dataset_list(self):
  40. return self.dataset_list