experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Merge pull request #44923 from Intel-tensorflow:tenglu/optimize_cnms

Browse Source 
PiperOrigin-RevId: 347022044
Change-Id: I899955358bf05d5d68f7061683a973608a056f31
This commit is contained in:
TensorFlower Gardener 2020-12-11 10:15:41 -08:00
commit 5fb1a09188
1 changed files with 4 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
tensorflow/core/kernels/image/non_max_suppression_op.cc
View File

@ -320,12 +320,6 @@ void DoNMSPerClass(int batch_idx, int class_idx, const float* boxes_data,
} }
} }
// Copy class_boxes_data to a tensor
TensorShape boxesShape({num_boxes, 4});
Tensor boxes(DT_FLOAT, boxesShape);
std::copy_n(class_boxes_data.begin(), class_boxes_data.size(),
boxes.unaligned_flat<float>().data());
// Do NMS, get the candidate indices of form vector<int> // Do NMS, get the candidate indices of form vector<int>
// Data structure for selection candidate in NMS. // Data structure for selection candidate in NMS.
struct Candidate { struct Candidate {
@ -347,9 +341,10 @@ void DoNMSPerClass(int batch_idx, int class_idx, const float* boxes_data,
Candidate next_candidate; Candidate next_candidate;
std::sort(candidate_vector.begin(), candidate_vector.end(), cmp); std::sort(candidate_vector.begin(), candidate_vector.end(), cmp);
const Tensor const_boxes = boxes; // Move class_boxes_data to a tensor
typename TTypes<float, 2>::ConstTensor boxes_data_t = Eigen::array<Eigen::DenseIndex, 2> boxesShape = {num_boxes, 4};
const_boxes.tensor<float, 2>(); typename TTypes<float, 2>::ConstTensor boxes_data_t(class_boxes_data.data(),
boxesShape);
int candidate_idx = 0; int candidate_idx = 0;
float iou; float iou;
while (selected.size() < size_per_class && while (selected.size() < size_per_class &&