| 53 | |
| 54 | template <typename Dtype> |
| 55 | void ArgMaxLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom, |
| 56 | const vector<Blob<Dtype>*>& top) { |
| 57 | const Dtype* bottom_data = bottom[0]->cpu_data(); |
| 58 | Dtype* top_data = top[0]->mutable_cpu_data(); |
| 59 | int dim, axis_dist; |
| 60 | if (has_axis_) { |
| 61 | dim = bottom[0]->shape(axis_); |
| 62 | // Distance between values of axis in blob |
| 63 | axis_dist = bottom[0]->count(axis_) / dim; |
| 64 | } else { |
| 65 | dim = bottom[0]->count(1); |
| 66 | axis_dist = 1; |
| 67 | } |
| 68 | int num = bottom[0]->count() / dim; |
| 69 | std::vector<std::pair<Dtype, int> > bottom_data_vector(dim); |
| 70 | for (int i = 0; i < num; ++i) { |
| 71 | for (int j = 0; j < dim; ++j) { |
| 72 | bottom_data_vector[j] = std::make_pair( |
| 73 | bottom_data[(i / axis_dist * dim + j) * axis_dist + i % axis_dist], j); |
| 74 | } |
| 75 | std::partial_sort( |
| 76 | bottom_data_vector.begin(), bottom_data_vector.begin() + top_k_, |
| 77 | bottom_data_vector.end(), std::greater<std::pair<Dtype, int> >()); |
| 78 | for (int j = 0; j < top_k_; ++j) { |
| 79 | if (out_max_val_) { |
| 80 | if (has_axis_) { |
| 81 | // Produces max_val per axis |
| 82 | top_data[(i / axis_dist * top_k_ + j) * axis_dist + i % axis_dist] |
| 83 | = bottom_data_vector[j].first; |
| 84 | } else { |
| 85 | // Produces max_ind and max_val |
| 86 | top_data[2 * i * top_k_ + j] = bottom_data_vector[j].second; |
| 87 | top_data[2 * i * top_k_ + top_k_ + j] = bottom_data_vector[j].first; |
| 88 | } |
| 89 | } else { |
| 90 | // Produces max_ind per axis |
| 91 | top_data[(i / axis_dist * top_k_ + j) * axis_dist + i % axis_dist] |
| 92 | = bottom_data_vector[j].second; |
| 93 | } |
| 94 | } |
| 95 | } |
| 96 | } |
| 97 | |
| 98 | INSTANTIATE_CLASS(ArgMaxLayer); |
| 99 | REGISTER_LAYER_CLASS(ArgMax); |