mix_sampler.py 3.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  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 absolute_import
  15. from __future__ import division
  16. from paddle.io import DistributedBatchSampler, Sampler
  17. from paddlex.ppcls.utils import logger
  18. from paddlex.ppcls.data.dataloader.mix_dataset import MixDataset
  19. from paddlex.ppcls.data import dataloader
  20. class MixSampler(DistributedBatchSampler):
  21. def __init__(self, dataset, batch_size, sample_configs, iter_per_epoch):
  22. super().__init__(dataset, batch_size)
  23. assert isinstance(dataset,
  24. MixDataset), "MixSampler only support MixDataset"
  25. self.sampler_list = []
  26. self.batch_size = batch_size
  27. self.start_list = []
  28. self.length = iter_per_epoch
  29. dataset_list = dataset.get_dataset_list()
  30. batch_size_left = self.batch_size
  31. self.iter_list = []
  32. for i, config_i in enumerate(sample_configs):
  33. self.start_list.append(dataset_list[i][1])
  34. sample_method = config_i.pop("name")
  35. ratio_i = config_i.pop("ratio")
  36. if i < len(sample_configs) - 1:
  37. batch_size_i = int(self.batch_size * ratio_i)
  38. batch_size_left -= batch_size_i
  39. else:
  40. batch_size_i = batch_size_left
  41. assert batch_size_i <= len(dataset_list[i][2])
  42. config_i["batch_size"] = batch_size_i
  43. if sample_method == "DistributedBatchSampler":
  44. sampler_i = DistributedBatchSampler(dataset_list[i][2],
  45. **config_i)
  46. else:
  47. sampler_i = getattr(dataloader, sample_method)(
  48. dataset_list[i][2], **config_i)
  49. self.sampler_list.append(sampler_i)
  50. self.iter_list.append(iter(sampler_i))
  51. self.length += len(dataset_list[i][2]) * ratio_i
  52. self.iter_counter = 0
  53. def __iter__(self):
  54. while self.iter_counter < self.length:
  55. batch = []
  56. for i, iter_i in enumerate(self.iter_list):
  57. batch_i = next(iter_i, None)
  58. if batch_i is None:
  59. iter_i = iter(self.sampler_list[i])
  60. self.iter_list[i] = iter_i
  61. batch_i = next(iter_i, None)
  62. assert batch_i is not None, "dataset {} return None".format(
  63. i)
  64. batch += [idx + self.start_list[i] for idx in batch_i]
  65. if len(batch) == self.batch_size:
  66. self.iter_counter += 1
  67. yield batch
  68. else:
  69. logger.info("Some dataset reaches end")
  70. self.iter_counter = 0
  71. def __len__(self):
  72. return self.length