Caffe2 - C++ API
A deep learning, cross platform ML framework
spatial_softmax_with_loss_op.cc
1 
17 #include "spatial_softmax_with_loss_op.h"
18 #include "softmax_shared.h"
19 
20 namespace caffe2 {
21 
22 REGISTER_CPU_OPERATOR(
23  SpatialSoftmaxWithLoss,
24  SpatialSoftmaxWithLossOp<float, CPUContext>);
25 REGISTER_CPU_OPERATOR(
26  SpatialSoftmaxWithLossGradient,
27  SpatialSoftmaxWithLossGradientOp<float, CPUContext>);
28 
29 // Input: X (logits), T (labels); Output: P (probs), Y
30 OPERATOR_SCHEMA(SpatialSoftmaxWithLoss)
31  .NumInputs(2, 3)
32  .NumOutputs(2)
33  .TensorInferenceFunction(
34  [](const OperatorDef& def, const vector<TensorShape>& in) {
35  ArgumentHelper helper(def);
36  vector<TensorShape> out(2);
37 
38  auto logits = in[0]; // Tensor with Shape [batch_size, num_classes]
39  auto labels = in[1]; // Tensor with shape [batch_size, ]
40  auto batch_size = logits.dims().Get(0);
41  auto num_classes = logits.dims().Get(1);
42 
43  CAFFE_ENFORCE_EQ(logits.dims_size(), 4);
44  CAFFE_ENFORCE_EQ(labels.dims_size(), 3);
45  out[0].set_data_type(logits.data_type());
46  out[0].add_dims(batch_size);
47  out[0].add_dims(num_classes);
48  out[0].add_dims(in[0].dims(2));
49  out[0].add_dims(in[0].dims(3));
50  // Output 2 is scalar shape, so no dims added
51  return out;
52  })
53  .SetDoc(R"DOC(
54 Combined Spatial Softmax and Cross-Entropy loss operator.
55 Similar to SoftmaxWithLoss, this operator computes the spatial softmax
56 normalized values for each layer in the batch of the given input, after which
57 cross-entropy loss is computed. This operator is numerically more stable than
58 separate Softmax and CrossEntropy ops. The inputs are a 2-D tensor
59 (Tensor<float>) of size (batch_size x input_feature_dimensions) and tensor of
60 labels (ground truth).
61 Output is tensor with the probability for each label in a pixel for each example
62 (N x D x W x H) and averaged loss (scalar).
63 For spatial softmax, weighting is by x,y position of the input.
64 )DOC")
65  .Input(0, "logits", "Unscaled log probabilities")
66  .Input(1, "labels", "Ground truth")
67  .Input(
68  2,
69  "weight_tensor",
70  "Optional blob to be used to weight the samples for the loss. With\
71  spatial set, weighting is by x,y of the input")
72  .Output(0, "softmax", "Tensor with softmax cross entropy loss")
73  .Output(1, "loss", "Average loss");
74 
75 // Input: X, T, P, dY; Output: dX
76 OPERATOR_SCHEMA(SpatialSoftmaxWithLossGradient).NumOutputs(1);
77 
78 #define DONT_CARE (-1)
79 
80 template <>
81 bool SpatialSoftmaxWithLossOp<float, CPUContext>::RunOnDevice() {
82  auto& X = Input(0); // Logits
83  auto& T = Input(1); // Labels / targets
84  auto* P = Output(0); // Probabilities from softmax
85  auto* avg_loss = Output(1); // Average loss
86  int N, D;
87  N = X.dim32(0);
88  D = X.dim32(1);
89  P->ResizeLike(X);
90 
91  if (sum_multiplier_.size() != D) {
92  sum_multiplier_.Resize(D);
93  math::Set<float, CPUContext>(
94  D, 1.f, sum_multiplier_.mutable_data<float>(), &context_);
95  }
96 
97  float* Pdata = P->mutable_data<float>();
98  const float* weights = (InputSize() > 2 ? Input(2).data<float>() : nullptr);
99  CAFFE_ENFORCE_EQ(X.ndim(), 4);
100  CAFFE_ENFORCE_EQ(T.ndim(), 3);
101  CAFFE_ENFORCE_EQ(T.dim32(0), N);
102 
103  int H = X.dim32(2);
104  int W = X.dim32(3);
105 
106  const float* Xdata = X.data<float>();
107 
108  for (int i = 0; i < N; ++i) {
109  for (int y = 0; y < H; ++y) {
110  for (int x = 0; x < W; ++x) {
111  // Subtract max on each cell for numerical reasons
112  float max_val = (-1e20f);
113  for (int c = 0; c < D; ++c) {
114  // TODO optimize
115  int idx = i * (H * W * D) + c * (H * W) + y * W + x;
116  max_val = std::max(max_val, Xdata[idx]);
117  }
118 
119  // Exponentiate
120  float expsum = 0.0f;
121  for (int c = 0; c < D; ++c) {
122  int idx = i * (H * W * D) + c * (H * W) + y * W + x;
123  float expx = exp(Xdata[idx] - max_val);
124  Pdata[idx] = expx;
125  expsum += expx;
126  }
127 
128  // Normalize
129  for (int c = 0; c < D; ++c) {
130  int idx = i * (H * W * D) + c * (H * W) + y * W + x;
131  Pdata[idx] /= expsum;
132  }
133  }
134  }
135  }
136 
137  // Compute the avg cross-entropy loss
138  avg_loss->Resize(vector<TIndex>());
139  float* avg_loss_data = avg_loss->mutable_data<float>();
140  const int* label_data = T.data<int>();
141 
142  float sum_label_xent = 0.0f;
143  float total_weight = 0.0;
144 
145  for (int y = 0; y < H; y++) {
146  for (int x = 0; x < W; x++) {
147  for (int i = 0; i < N; i++) {
148  int label_idx = i * H * W + y * W + x;
149  int label = label_data[label_idx];
150  if (label != DONT_CARE) {
151  CAFFE_ENFORCE(
152  label < D && label >= 0,
153  "Label seems incorrect:label value larger than number of classes",
154  label_data[i],
155  " vs ",
156  D);
157  int idx = i * (H * W * D) + label * (H * W) + y * W + x;
158  float w = weights ? weights[label_idx] : 1.0;
159  total_weight += w;
160  sum_label_xent += -log(std::max(Pdata[idx], 1e-20f)) * w;
161  }
162  }
163  }
164  }
165  if (total_weight != 0.0) {
166  *avg_loss_data = sum_label_xent / total_weight;
167  } else {
168  *avg_loss_data = 0.0;
169  }
170  return true;
171 }
172 
173 template <>
174 bool SpatialSoftmaxWithLossGradientOp<float, CPUContext>::RunOnDevice() {
175  auto& X = Input(0); // Logits
176  auto& T = Input(1); // Labels / targets
177  // Input(2) is weights if given
178  auto& P = Input(InputSize() - 2); // Probabilities from softmax
179  auto& d_avg_loss = Input(InputSize() - 1); // Gradient w.r.t. avg loss
180  auto* dX = Output(0);
181  const float* weights = (InputSize() > 4 ? Input(2).data<float>() : nullptr);
182  int N, D;
183  N = X.dim32(0);
184  D = X.dim32(1);
185  dX->ResizeLike(X);
186  CAFFE_ENFORCE_EQ(T.dim32(0), N);
187  CAFFE_ENFORCE_EQ(X.ndim(), 4);
188  CAFFE_ENFORCE_EQ(T.ndim(), 3);
189 
190  int H = X.dim32(2);
191  int W = X.dim32(3);
192 
193  const float* Pdata = P.data<float>();
194  float* dX_data = dX->mutable_data<float>();
195  const int* label_data = T.data<int>();
196 
197  // Copy softmax probabilities into dX. All but the neuron
198  // corresponding to the correct label has gradient equaling e(x_j)
199  // which is the probability under softmax.
200  context_.Copy<float, CPUContext, CPUContext>(P.size(), Pdata, dX_data);
201 
202  float total_weight = 0.0f;
203  for (int y = 0; y < H; ++y) {
204  for (int x = 0; x < W; ++x) {
205  for (int i = 0; i < N; ++i) {
206  int label_idx = i * H * W + y * W + x;
207  int label = label_data[label_idx];
208 
209  if (label != DONT_CARE) {
210  int idx = i * (H * W * D) + label * (H * W) + y * W + x;
211 
212  dX_data[idx] = (dX_data[idx] - 1.0);
213 
214  if (weights != nullptr) {
215  float weight = weights[label_idx];
216  for (int c = 0; c < D; ++c) {
217  int k = i * (H * W * D) + c * (H * W) + y * W + x;
218  dX_data[k] *= weight;
219  }
220  total_weight += weight;
221  } else {
222  total_weight += 1.0;
223  }
224  } else {
225  // Set gradient to zero for coordinates where we have dont care
226  for (int c = 0; c < D; ++c) {
227  int idx = i * (H * W * D) + c * (H * W) + y * W + x;
228  dX_data[idx] = 0;
229  }
230  }
231  }
232  }
233  }
234 
235  if (total_weight > 0) {
236  math::Scale<float, CPUContext>(
237  dX->size(),
238  scale_ / total_weight,
239  dX->data<float>(),
240  dX_data,
241  &context_);
242  }
243  math::Scale<float, CPUContext>(
244  dX->size(),
245  d_avg_loss.data<float>(),
246  dX->data<float>(),
247  dX->mutable_data<float>(),
248  &context_);
249  return true;
250 }
251 
252 namespace {
253 class GetSoftmaxWithLossGradient : public GradientMakerBase {
254  using GradientMakerBase::GradientMakerBase;
255  vector<OperatorDef> GetGradientDefs() override {
256  vector<string> blob_names{
257  {I(0), I(1), O(0), GO(1)},
258  };
259 
260  // Add weight blob, if given
261  if (def_.input_size() == 3) {
262  blob_names.emplace(blob_names.begin() + 2, I(2));
263  }
264  return SingleGradientDef(
265  "SpatialSoftmaxWithLossGradient",
266  "",
267  blob_names,
268  vector<string>{GI(0)});
269  }
270 };
271 
272 REGISTER_GRADIENT(SpatialSoftmaxWithLoss, GetSoftmaxWithLossGradient);
273 }
274 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.