caffe2 Softmax, SoftmaxWithLoss 使用

本文主要记录caffe2的Softmax与SoftmaxWithLoss op 如何使用(输入格式), 特别介绍如何应用在FCN的语义分割中.
1 Softmax
输入:2D, 但可以不是,计算时代码会强制转换
输出: 与输入相同的维度
arg:axis 默认等于1, 其实没太搞懂
在FCN的输出probs时以下操作, 直接上测试的代码

from caffe2.python import core, workspace, cnn
import numpy as np
#fetch blob
np.random.seed(0)
X0 = np.random.rand(2,2,2,2) * 10 # NCHW, C equal n_classes
label = np.array([[[0, 0], [1, 1]],[[0, 0], [1, 1]]], dtype=np.int32) # NHW, 3D

workspace.FeedBlob('X0', np.array(X0, dtype=np.float32))
workspace.FeedBlob('label', np.array(label, dtype=np.int32)) # label is int type

model = cnn.CNNModelHelper()
X0_NHWC = model.net.NCHW2NHWC('X0', 'X0_NHWC') # NCHW -> NHWC
model.net.Softmax('X0_NHWC', 'probs_NHWC', axis=3) # using when test/val
model.net.NHWC2NCHW('probs_NHWC', 'probs_NCHW') 

2 SoftmaxWithLoss op
输入:
logits: 必须时2D(N*D)
label: is either a 1D array of size N (batch size), or a 2D array of size N x 1 (batch size). Each entry in the label vector indicates which is the correct class; as such, each entry must be between 0 and D - 1, inclusive, where D is the total number of classes.
tip: you need a reshape on your logits and label

X0_NHWC = model.net.NCHW2NHWC('X0', 'X0_NHWC')
model.Reshape('X0_NHWC',['X0_reshape', 'old_shape'], shape=[-1, 2]) # (8, 2), n_classes: 2
model.Reshape('label', ['label_reshape', 'labelOldShape'], shape=[-1,]) # shape: (8,)

model.net.SoftmaxWithLoss(['X0_reshape', 'label_reshape'], ['probs', 'loss'])