does caffe2 translator handle the axis parameter in SoftMax?

[alekseic]

I have a simple CNN. When I send data as a = np.random.rand(1,2,3,3), I want to obtain one output of shape (1, 9, 2), in which the SoftMax is taken along axis = -1. So, in this way np.sum(out) = 9, and for each index i, I expect to have out[0][i][0] + out[0][i][1] = 1. And it works in my Caffe model. But no matter what "axis" I specify in SoftMax parameters, when I translate this net to Caffe2, I have np.sum(out) = 1, which means that the SoftMax is taken for all elements, not for pairs. Could it be an issue with caffe translator? The CNN is below:

name: "MyNet" input: "data" input_dim: 1 input_dim: 2 input_dim: 3 input_dim: 3

layer { name: "reshape" type: "Reshape" bottom: "data" top: "reshape" reshape_param { shape: {dim: 0 dim: -1 dim: 2} } # shape becomes (1, 9, 2) }

layer { name: "prob1" type: "Softmax" bottom: "reshape" top: "prob1" softmax_param { axis : -1 } }

UPD: It seems that the axis parameter in SoftMax doesn't exist at all in Caffe2; and Caffe2 doesn't work with (more than 2)-dimensional tensors.

Had to replace my reshape layer: reshape_param { shape: {dim: -1 dim: 2} } In this way, a 2D tensor is passed to SoftMax, and it works in the way I wanted.

[paghdv]

I wonder if anyone figure this out, the solution proposed here is more of a hack. Softmax axis are required for tasks like semantic image segmentation.