DistributedRandomIdentitySampler.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990
  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 copy
  19. import random
  20. from paddle.io import DistributedBatchSampler, Sampler
  21. class DistributedRandomIdentitySampler(DistributedBatchSampler):
  22. """
  23. Randomly sample N identities, then for each identity,
  24. randomly sample K instances, therefore batch size is N*K.
  25. Args:
  26. - data_source (list): list of (img_path, pid, camid).
  27. - num_instances (int): number of instances per identity in a batch.
  28. - batch_size (int): number of examples in a batch.
  29. """
  30. def __init__(self, dataset, batch_size, num_instances, drop_last, **args):
  31. self.dataset = dataset
  32. self.batch_size = batch_size
  33. self.num_instances = num_instances
  34. self.drop_last = drop_last
  35. self.num_pids_per_batch = self.batch_size // self.num_instances
  36. self.index_dic = defaultdict(list)
  37. for index, pid in enumerate(self.dataset.labels):
  38. self.index_dic[pid].append(index)
  39. self.pids = list(self.index_dic.keys())
  40. # estimate number of examples in an epoch
  41. self.length = 0
  42. for pid in self.pids:
  43. idxs = self.index_dic[pid]
  44. num = len(idxs)
  45. if num < self.num_instances:
  46. num = self.num_instances
  47. self.length += num - num % self.num_instances
  48. def __iter__(self):
  49. batch_idxs_dict = defaultdict(list)
  50. for pid in self.pids:
  51. idxs = copy.deepcopy(self.index_dic[pid])
  52. if len(idxs) < self.num_instances:
  53. idxs = np.random.choice(
  54. idxs, size=self.num_instances, replace=True)
  55. random.shuffle(idxs)
  56. batch_idxs = []
  57. for idx in idxs:
  58. batch_idxs.append(idx)
  59. if len(batch_idxs) == self.num_instances:
  60. batch_idxs_dict[pid].append(batch_idxs)
  61. batch_idxs = []
  62. avai_pids = copy.deepcopy(self.pids)
  63. final_idxs = []
  64. while len(avai_pids) >= self.num_pids_per_batch:
  65. selected_pids = random.sample(avai_pids, self.num_pids_per_batch)
  66. for pid in selected_pids:
  67. batch_idxs = batch_idxs_dict[pid].pop(0)
  68. final_idxs.extend(batch_idxs)
  69. if len(batch_idxs_dict[pid]) == 0:
  70. avai_pids.remove(pid)
  71. _sample_iter = iter(final_idxs)
  72. batch_indices = []
  73. for idx in _sample_iter:
  74. batch_indices.append(idx)
  75. if len(batch_indices) == self.batch_size:
  76. yield batch_indices
  77. batch_indices = []
  78. if not self.drop_last and len(batch_indices) > 0:
  79. yield batch_indices
  80. def __len__(self):
  81. if self.drop_last:
  82. return self.length // self.batch_size
  83. else:
  84. return (self.length + self.batch_size - 1) // self.batch_size