retrieval.py 6.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169
  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 __future__ import print_function
  17. import platform
  18. import paddle
  19. from paddlex.ppcls.utils import logger
  20. def retrieval_eval(engine, epoch_id=0):
  21. engine.model.eval()
  22. # step1. build gallery
  23. if engine.gallery_query_dataloader is not None:
  24. gallery_feas, gallery_img_id, gallery_unique_id = cal_feature(
  25. engine, name='gallery_query')
  26. query_feas, query_img_id, query_query_id = gallery_feas, gallery_img_id, gallery_unique_id
  27. else:
  28. gallery_feas, gallery_img_id, gallery_unique_id = cal_feature(
  29. engine, name='gallery')
  30. query_feas, query_img_id, query_query_id = cal_feature(
  31. engine, name='query')
  32. # step2. do evaluation
  33. sim_block_size = engine.config["Global"].get("sim_block_size", 64)
  34. sections = [sim_block_size] * (len(query_feas) // sim_block_size)
  35. if len(query_feas) % sim_block_size:
  36. sections.append(len(query_feas) % sim_block_size)
  37. fea_blocks = paddle.split(query_feas, num_or_sections=sections)
  38. if query_query_id is not None:
  39. query_id_blocks = paddle.split(
  40. query_query_id, num_or_sections=sections)
  41. image_id_blocks = paddle.split(query_img_id, num_or_sections=sections)
  42. metric_key = None
  43. if engine.eval_loss_func is None:
  44. metric_dict = {metric_key: 0.}
  45. else:
  46. metric_dict = dict()
  47. for block_idx, block_fea in enumerate(fea_blocks):
  48. similarity_matrix = paddle.matmul(
  49. block_fea, gallery_feas, transpose_y=True)
  50. if query_query_id is not None:
  51. query_id_block = query_id_blocks[block_idx]
  52. query_id_mask = (query_id_block != gallery_unique_id.t())
  53. image_id_block = image_id_blocks[block_idx]
  54. image_id_mask = (image_id_block != gallery_img_id.t())
  55. keep_mask = paddle.logical_or(query_id_mask, image_id_mask)
  56. similarity_matrix = similarity_matrix * keep_mask.astype(
  57. "float32")
  58. else:
  59. keep_mask = None
  60. metric_tmp = engine.eval_metric_func(similarity_matrix,
  61. image_id_blocks[block_idx],
  62. gallery_img_id, keep_mask)
  63. for key in metric_tmp:
  64. if key not in metric_dict:
  65. metric_dict[key] = metric_tmp[key] * block_fea.shape[
  66. 0] / len(query_feas)
  67. else:
  68. metric_dict[key] += metric_tmp[key] * block_fea.shape[
  69. 0] / len(query_feas)
  70. metric_info_list = []
  71. for key in metric_dict:
  72. if metric_key is None:
  73. metric_key = key
  74. metric_info_list.append("{}: {:.5f}".format(key, metric_dict[key]))
  75. metric_msg = ", ".join(metric_info_list)
  76. logger.info("[Eval][Epoch {}][Avg]{}".format(epoch_id, metric_msg))
  77. return metric_dict[metric_key]
  78. def cal_feature(engine, name='gallery'):
  79. all_feas = None
  80. all_image_id = None
  81. all_unique_id = None
  82. has_unique_id = False
  83. if name == 'gallery':
  84. dataloader = engine.gallery_dataloader
  85. elif name == 'query':
  86. dataloader = engine.query_dataloader
  87. elif name == 'gallery_query':
  88. dataloader = engine.gallery_query_dataloader
  89. else:
  90. raise RuntimeError("Only support gallery or query dataset")
  91. max_iter = len(dataloader) - 1 if platform.system() == "Windows" else len(
  92. dataloader)
  93. for idx, batch in enumerate(dataloader): # load is very time-consuming
  94. if idx >= max_iter:
  95. break
  96. if idx % engine.config["Global"]["print_batch_step"] == 0:
  97. logger.info(
  98. f"{name} feature calculation process: [{idx}/{len(dataloader)}]"
  99. )
  100. if engine.use_dali:
  101. batch = [
  102. paddle.to_tensor(batch[0]['data']),
  103. paddle.to_tensor(batch[0]['label'])
  104. ]
  105. batch = [paddle.to_tensor(x) for x in batch]
  106. batch[1] = batch[1].reshape([-1, 1]).astype("int64")
  107. if len(batch) == 3:
  108. has_unique_id = True
  109. batch[2] = batch[2].reshape([-1, 1]).astype("int64")
  110. out = engine.model(batch[0], batch[1])
  111. batch_feas = out["features"]
  112. # do norm
  113. if engine.config["Global"].get("feature_normalize", True):
  114. feas_norm = paddle.sqrt(
  115. paddle.sum(paddle.square(batch_feas), axis=1, keepdim=True))
  116. batch_feas = paddle.divide(batch_feas, feas_norm)
  117. # do binarize
  118. if engine.config["Global"].get("feature_binarize") == "round":
  119. batch_feas = paddle.round(batch_feas).astype("float32") * 2.0 - 1.0
  120. if engine.config["Global"].get("feature_binarize") == "sign":
  121. batch_feas = paddle.sign(batch_feas).astype("float32")
  122. if all_feas is None:
  123. all_feas = batch_feas
  124. if has_unique_id:
  125. all_unique_id = batch[2]
  126. all_image_id = batch[1]
  127. else:
  128. all_feas = paddle.concat([all_feas, batch_feas])
  129. all_image_id = paddle.concat([all_image_id, batch[1]])
  130. if has_unique_id:
  131. all_unique_id = paddle.concat([all_unique_id, batch[2]])
  132. if engine.use_dali:
  133. dataloader.reset()
  134. if paddle.distributed.get_world_size() > 1:
  135. feat_list = []
  136. img_id_list = []
  137. unique_id_list = []
  138. paddle.distributed.all_gather(feat_list, all_feas)
  139. paddle.distributed.all_gather(img_id_list, all_image_id)
  140. all_feas = paddle.concat(feat_list, axis=0)
  141. all_image_id = paddle.concat(img_id_list, axis=0)
  142. if has_unique_id:
  143. paddle.distributed.all_gather(unique_id_list, all_unique_id)
  144. all_unique_id = paddle.concat(unique_id_list, axis=0)
  145. logger.info("Build {} done, all feat shape: {}, begin to eval..".format(
  146. name, all_feas.shape))
  147. return all_feas, all_image_id, all_unique_id