| 31 | |
| 32 | template <typename Dtype> |
| 33 | void ArgMaxLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom, |
| 34 | const vector<Blob<Dtype>*>& top) { |
| 35 | int num_top_axes = bottom[0]->num_axes(); |
| 36 | if ( num_top_axes < 3 ) num_top_axes = 3; |
| 37 | std::vector<int> shape(num_top_axes, 1); |
| 38 | if (has_axis_) { |
| 39 | // Produces max_ind or max_val per axis |
| 40 | shape = bottom[0]->shape(); |
| 41 | shape[axis_] = top_k_; |
| 42 | } else { |
| 43 | shape[0] = bottom[0]->shape(0); |
| 44 | // Produces max_ind |
| 45 | shape[2] = top_k_; |
| 46 | if (out_max_val_) { |
| 47 | // Produces max_ind and max_val |
| 48 | shape[1] = 2; |
| 49 | } |
| 50 | } |
| 51 | top[0]->Reshape(shape); |
| 52 | } |
| 53 | |
| 54 | template <typename Dtype> |
| 55 | void ArgMaxLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom, |