pk_sampler.py 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105
  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 collections import defaultdict
  17. import numpy as np
  18. import random
  19. from paddle.io import DistributedBatchSampler
  20. from paddlex.ppcls.utils import logger
  21. class PKSampler(DistributedBatchSampler):
  22. """
  23. First, randomly sample P identities.
  24. Then for each identity randomly sample K instances.
  25. Therefore batch size is P*K, and the sampler called PKSampler.
  26. Args:
  27. dataset (paddle.io.Dataset): list of (img_path, pid, cam_id).
  28. sample_per_id(int): number of instances per identity in a batch.
  29. batch_size (int): number of examples in a batch.
  30. shuffle(bool): whether to shuffle indices order before generating
  31. batch indices. Default False.
  32. """
  33. def __init__(self,
  34. dataset,
  35. batch_size,
  36. sample_per_id,
  37. shuffle=True,
  38. drop_last=True,
  39. sample_method="sample_avg_prob"):
  40. super().__init__(
  41. dataset, batch_size, shuffle=shuffle, drop_last=drop_last)
  42. assert batch_size % sample_per_id == 0, \
  43. "PKSampler configs error, Sample_per_id must be a divisor of batch_size."
  44. assert hasattr(self.dataset,
  45. "labels"), "Dataset must have labels attribute."
  46. self.sample_per_label = sample_per_id
  47. self.label_dict = defaultdict(list)
  48. self.sample_method = sample_method
  49. for idx, label in enumerate(self.dataset.labels):
  50. self.label_dict[label].append(idx)
  51. self.label_list = list(self.label_dict)
  52. assert len(self.label_list) * self.sample_per_label > self.batch_size, \
  53. "batch size should be smaller than "
  54. if self.sample_method == "id_avg_prob":
  55. self.prob_list = np.array([1 / len(self.label_list)] *
  56. len(self.label_list))
  57. elif self.sample_method == "sample_avg_prob":
  58. counter = []
  59. for label_i in self.label_list:
  60. counter.append(len(self.label_dict[label_i]))
  61. self.prob_list = np.array(counter) / sum(counter)
  62. else:
  63. logger.error(
  64. "PKSampler only support id_avg_prob and sample_avg_prob sample method, "
  65. "but receive {}.".format(self.sample_method))
  66. diff = np.abs(sum(self.prob_list) - 1)
  67. if diff > 0.00000001:
  68. self.prob_list[-1] = 1 - sum(self.prob_list[:-1])
  69. if self.prob_list[-1] > 1 or self.prob_list[-1] < 0:
  70. logger.error("PKSampler prob list error")
  71. else:
  72. logger.info(
  73. "PKSampler: sum of prob list not equal to 1, diff is {}, change the last prob".
  74. format(diff))
  75. def __iter__(self):
  76. label_per_batch = self.batch_size // self.sample_per_label
  77. for _ in range(len(self)):
  78. batch_index = []
  79. batch_label_list = np.random.choice(
  80. self.label_list,
  81. size=label_per_batch,
  82. replace=False,
  83. p=self.prob_list)
  84. for label_i in batch_label_list:
  85. label_i_indexes = self.label_dict[label_i]
  86. if self.sample_per_label <= len(label_i_indexes):
  87. batch_index.extend(
  88. np.random.choice(
  89. label_i_indexes,
  90. size=self.sample_per_label,
  91. replace=False))
  92. else:
  93. batch_index.extend(
  94. np.random.choice(
  95. label_i_indexes,
  96. size=self.sample_per_label,
  97. replace=True))
  98. if not self.drop_last or len(batch_index) == self.batch_size:
  99. yield batch_index