背景#
似乎没什么背景，继续看 caffe 代码。
argmax 的作用是返回一个 blob 某个维度（或者 batch_size 之后的维度）的 top_k 的 index（或者 pair(index, value)）。
proto#
还是先看 proto。
1
2message ArgMaxParameter {
3 // If true produce pairs (argmax, maxval)
4 optional bool out_max_val = 1 [default = false];
5 optional uint32 top_k = 2 [default = 1];
6 // The axis along which to maximise -- may be negative to index from the
7 // end (e.g., -1 for the last axis).
8 // By default ArgMaxLayer maximizes over the flattened trailing dimensions
9 // for each index of the first / num dimension.
10 optional int32 axis = 3;
11}out_max_val 为真表示输出 (index, val) 的 pair，否则只输出 index（存疑）。
top_k 应该是要取最大的 top k 个元素。
axis 是要求最大的维度，默认情况是把 batch_size 之后的维度 flatten 之后求 argmax。
c++ 实现#
先看 Reshape 的部分。
1
2template <typename Dtype>
3void ArgMaxLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
4 const vector<Blob<Dtype>*>& top) {
5 int num_top_axes = bottom[0]->num_axes();
6 if ( num_top_axes < 3 ) num_top_axes = 3;
7 std::vector<int> shape(num_top_axes, 1);
8 if (has_axis_) {
9 // Produces max_ind or max_val per axis
10 shape = bottom[0]->shape();
11 shape[axis_] = top_k_;
12 // axis非默认参数的case: 只有求max的那个维度会变,其他都不变
13 // 问题: out_max_val似乎只适用在axis为默认参数的情况?
14 } else {
15 shape[0] = bottom[0]->shape(0);
16 // Produces max_ind
17 shape[2] = top_k_;
18 // 不是只拿到第top_k,而是拿到top_k的k个结果
19 if (out_max_val_) {
20 // Produces max_ind and max_val
21 shape[1] = 2;
22 }
23 // 默认axis参数得到的top blob的shape 为(batch_size,1或者2,top_k)
24 // 因为会把batch后面的维度flatten 然后求max
25 }
26 top[0]->Reshape(shape);
27}添加了一些注释。有一个疑问是，axis 和 out_max_val_ 这两个参数似乎不支持同时处理。
继续看 forward。
详细代码
1
2template <typename Dtype>
3void ArgMaxLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
4 const vector<Blob<Dtype>*>& top) {
5 const Dtype* bottom_data = bottom[0]->cpu_data();
6 Dtype* top_data = top[0]->mutable_cpu_data();
7 int dim, axis_dist;
8 if (has_axis_) {
9 dim = bottom[0]->shape(axis_);
10 // dim表示做argmax的维度一共有多少个值
11 // Distance between values of axis in blob
12 axis_dist = bottom[0]->count(axis_) / dim;
13 // 因为可能不在最末尾的维度做argmax,因此值在内存中未必是连续的
14 } else {
15 dim = bottom[0]->count(1);
16 // 从batch_size之后的维度数到最后
17 axis_dist = 1;
18 // 把末尾的几个维度flatten之后做argmax,在内存上这些值是连续的,因此axis_dist是1
19 }
20 int num = bottom[0]->count() / dim;
21 std::vector<std::pair<Dtype, int> > bottom_data_vector(dim);
22 for (int i = 0; i < num; ++i) {
23 for (int j = 0; j < dim; ++j) {
24 bottom_data_vector[j] = std::make_pair(
25 bottom_data[(i / axis_dist * dim + j) * axis_dist + i % axis_dist], j);
26 }
27 // 通过axis_dist控制,把要做argmax的元素从内存中不连续的位置传到一个连续的vector中,目的是做sort
28 std::partial_sort(
29 bottom_data_vector.begin(), bottom_data_vector.begin() + top_k_,
30 bottom_data_vector.end(), std::greater<std::pair<Dtype, int> >());
31 // 使得前top_k是最大的top_k个元素,后面的元素顺序任意
32 for (int j = 0; j < top_k_; ++j) {
33 if (out_max_val_) {
34 if (has_axis_) {
35 // Produces max_val per axis
36 top_data[(i / axis_dist * top_k_ + j) * axis_dist + i % axis_dist]
37 = bottom_data_vector[j].first;
38 //这个地方感觉有点问题... 就算是有axis参数不支持out_max_val... 输出的不也应该是index吗?
39 } else {
40 // Produces max_ind and max_val
41 top_data[2 * i * top_k_ + j] = bottom_data_vector[j].second;
42 top_data[2 * i * top_k_ + top_k_ + j] = bottom_data_vector[j].first;
43 }
44 } else {
45 // Produces max_ind per axis
46 top_data[(i / axis_dist * top_k_ + j) * axis_dist + i % axis_dist]
47 = bottom_data_vector[j].second;
48 }
49 }
50 }
51}由于可能不在末尾的维度求 max，因此求 max 的值可能在内存上是不连续的， 注意看 axis_dist 这个变量，表示的就是要求 argmax 的相邻元素在内存中的距离。
然后接下来的代码有些让人困惑… 即使不同时支持 out_max_val 和 axis 这两个参数、只有一个输出，那么输出的不也应该是 index 吗？
这个输出好像不是很对啊？去官方的 caffe 确认了一下，也是这样写的。
不是很确定这是不是预期的行为。
update:
看来不止我发现了这个问题 some doubts about argmax_layer (may be bug)，可惜 caffe 看起来已经没人维护了 2333。