/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserve. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. */ #pragma once #include #include #include "paddle/fluid/framework/op_registry.h" namespace paddle { namespace operators { using Tensor = framework::Tensor; template using EigenTensor = framework::EigenTensor; template using EigenVector = framework::EigenVector; using Array5 = Eigen::DSizes; template static inline bool isZero(T x) { return fabs(x) < 1e-6; } template static inline T sigmoid(T x) { return 1.0 / (exp(-1.0 * x) + 1.0); } template static inline T CalcMaskPointNum(const Tensor& mask) { auto mask_t = EigenVector::Flatten(mask); T count = 0.0; for (int i = 0; i < mask_t.dimensions()[0]; i++) { if (mask_t(i)) { count += 1.0; } } return count; } template static inline T CalcMSEWithMask(const Tensor& x, const Tensor& y, const Tensor& mask) { auto x_t = EigenVector::Flatten(x); auto y_t = EigenVector::Flatten(y); auto mask_t = EigenVector::Flatten(mask); T error_sum = 0.0; T points = 0.0; for (int i = 0; i < x_t.dimensions()[0]; i++) { if (mask_t(i)) { error_sum += pow(x_t(i) - y_t(i), 2); points += 1; } } return (error_sum / points); } template static void CalcMSEGradWithMask(Tensor* grad, const Tensor& x, const Tensor& y, const Tensor& mask, T mf) { auto grad_t = EigenVector::Flatten(*grad).setConstant(0.0); auto x_t = EigenVector::Flatten(x); auto y_t = EigenVector::Flatten(y); auto mask_t = EigenVector::Flatten(mask); for (int i = 0; i < x_t.dimensions()[0]; i++) { if (mask_t(i)) { grad_t(i) = 2.0 * (x_t(i) - y_t(i)) / mf; } } } template static inline T CalcBCEWithMask(const Tensor& x, const Tensor& y, const Tensor& mask) { auto x_t = EigenVector::Flatten(x); auto y_t = EigenVector::Flatten(y); auto mask_t = EigenVector::Flatten(mask); T error_sum = 0.0; T points = 0.0; for (int i = 0; i < x_t.dimensions()[0]; i++) { if (mask_t(i)) { error_sum += -1.0 * (y_t(i) * log(x_t(i)) + (1.0 - y_t(i)) * log(1.0 - x_t(i))); points += 1; } } return (error_sum / points); } template static inline void CalcBCEGradWithMask(Tensor* grad, const Tensor& x, const Tensor& y, const Tensor& mask, T mf) { auto grad_t = EigenVector::Flatten(*grad).setConstant(0.0); auto x_t = EigenVector::Flatten(x); auto y_t = EigenVector::Flatten(y); auto mask_t = EigenVector::Flatten(mask); for (int i = 0; i < x_t.dimensions()[0]; i++) { if (mask_t(i)) { grad_t(i) = ((1.0 - y_t(i)) / (1.0 - x_t(i)) - y_t(i) / x_t(i)) / mf; } } } template static void CalcPredResult(const Tensor& input, Tensor* pred_conf, Tensor* pred_class, Tensor* pred_x, Tensor* pred_y, Tensor* pred_w, Tensor* pred_h, const int anchor_num, const int class_num) { const int n = input.dims()[0]; const int h = input.dims()[2]; const int w = input.dims()[3]; const int box_attr_num = 5 + class_num; auto input_t = EigenTensor::From(input); auto pred_conf_t = EigenTensor::From(*pred_conf); auto pred_class_t = EigenTensor::From(*pred_class); auto pred_x_t = EigenTensor::From(*pred_x); auto pred_y_t = EigenTensor::From(*pred_y); auto pred_w_t = EigenTensor::From(*pred_w); auto pred_h_t = EigenTensor::From(*pred_h); for (int i = 0; i < n; i++) { for (int an_idx = 0; an_idx < anchor_num; an_idx++) { for (int j = 0; j < h; j++) { for (int k = 0; k < w; k++) { pred_x_t(i, an_idx, j, k) = sigmoid(input_t(i, box_attr_num * an_idx, j, k)); pred_y_t(i, an_idx, j, k) = sigmoid(input_t(i, box_attr_num * an_idx + 1, j, k)); pred_w_t(i, an_idx, j, k) = input_t(i, box_attr_num * an_idx + 2, j, k); pred_h_t(i, an_idx, j, k) = input_t(i, box_attr_num * an_idx + 3, j, k); pred_conf_t(i, an_idx, j, k) = sigmoid(input_t(i, box_attr_num * an_idx + 4, j, k)); for (int c = 0; c < class_num; c++) { pred_class_t(i, an_idx, j, k, c) = sigmoid(input_t(i, box_attr_num * an_idx + 5 + c, j, k)); } } } } } } template static T CalcBoxIoU(std::vector box1, std::vector box2) { T b1_x1 = box1[0] - box1[2] / 2; T b1_x2 = box1[0] + box1[2] / 2; T b1_y1 = box1[1] - box1[3] / 2; T b1_y2 = box1[1] + box1[3] / 2; T b2_x1 = box2[0] - box2[2] / 2; T b2_x2 = box2[0] + box2[2] / 2; T b2_y1 = box2[1] - box2[3] / 2; T b2_y2 = box2[1] + box2[3] / 2; T b1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1); T b2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1); T inter_rect_x1 = std::max(b1_x1, b2_x1); T inter_rect_y1 = std::max(b1_y1, b2_y1); T inter_rect_x2 = std::min(b1_x2, b2_x2); T inter_rect_y2 = std::min(b1_y2, b2_y2); T inter_area = std::max(inter_rect_x2 - inter_rect_x1, static_cast(0.0)) * std::max(inter_rect_y2 - inter_rect_y1, static_cast(0.0)); return inter_area / (b1_area + b2_area - inter_area); } template static void PreProcessGTBox(const Tensor& gt_boxes, const float ignore_thresh, std::vector anchors, const int grid_size, Tensor* obj_mask, Tensor* noobj_mask, Tensor* tx, Tensor* ty, Tensor* tw, Tensor* th, Tensor* tconf, Tensor* tclass) { const int n = gt_boxes.dims()[0]; const int b = gt_boxes.dims()[1]; const int anchor_num = anchors.size() / 2; auto gt_boxes_t = EigenTensor::From(gt_boxes); auto obj_mask_t = EigenTensor::From(*obj_mask).setConstant(0); auto noobj_mask_t = EigenTensor::From(*noobj_mask).setConstant(1); auto tx_t = EigenTensor::From(*tx).setConstant(0.0); auto ty_t = EigenTensor::From(*ty).setConstant(0.0); auto tw_t = EigenTensor::From(*tw).setConstant(0.0); auto th_t = EigenTensor::From(*th).setConstant(0.0); auto tconf_t = EigenTensor::From(*tconf).setConstant(0.0); auto tclass_t = EigenTensor::From(*tclass).setConstant(0.0); for (int i = 0; i < n; i++) { for (int j = 0; j < b; j++) { if (isZero(gt_boxes_t(i, j, 0)) && isZero(gt_boxes_t(i, j, 1)) && isZero(gt_boxes_t(i, j, 2)) && isZero(gt_boxes_t(i, j, 3)) && isZero(gt_boxes_t(i, j, 4))) { continue; } int gt_label = static_cast(gt_boxes_t(i, j, 0)); T gx = gt_boxes_t(i, j, 1) * grid_size; T gy = gt_boxes_t(i, j, 2) * grid_size; T gw = gt_boxes_t(i, j, 3) * grid_size; T gh = gt_boxes_t(i, j, 4) * grid_size; int gi = static_cast(gx); int gj = static_cast(gy); T max_iou = static_cast(0); T iou; int best_an_index = -1; std::vector gt_box({0, 0, gw, gh}); for (int an_idx = 0; an_idx < anchor_num; an_idx++) { std::vector anchor_shape({0, 0, static_cast(anchors[2 * an_idx]), static_cast(anchors[2 * an_idx + 1])}); iou = CalcBoxIoU(gt_box, anchor_shape); if (iou > max_iou) { max_iou = iou; best_an_index = an_idx; } if (iou > ignore_thresh) { noobj_mask_t(i, an_idx, gj, gi) = 0; } } obj_mask_t(i, best_an_index, gj, gi) = 1; noobj_mask_t(i, best_an_index, gj, gi) = 0; tx_t(i, best_an_index, gj, gi) = gx - gi; ty_t(i, best_an_index, gj, gi) = gy - gj; tw_t(i, best_an_index, gj, gi) = log(gw / anchors[2 * best_an_index]); th_t(i, best_an_index, gj, gi) = log(gh / anchors[2 * best_an_index + 1]); tclass_t(i, best_an_index, gj, gi, gt_label) = 1; tconf_t(i, best_an_index, gj, gi) = 1; } } } static void ExpandObjMaskByClassNum(Tensor* obj_mask_expand, const Tensor& obj_mask) { const int n = obj_mask_expand->dims()[0]; const int an_num = obj_mask_expand->dims()[1]; const int h = obj_mask_expand->dims()[2]; const int w = obj_mask_expand->dims()[3]; const int class_num = obj_mask_expand->dims()[4]; auto obj_mask_expand_t = EigenTensor::From(*obj_mask_expand); auto obj_mask_t = EigenTensor::From(obj_mask); obj_mask_expand_t = obj_mask_t.reshape(Array5(n, an_num, h, w, 1)) .broadcast(Array5(1, 1, 1, 1, class_num)); } template static void AddAllGradToInputGrad( Tensor* grad, T loss, const Tensor& pred_x, const Tensor& pred_y, const Tensor& pred_conf, const Tensor& pred_class, const Tensor& grad_x, const Tensor& grad_y, const Tensor& grad_w, const Tensor& grad_h, const Tensor& grad_conf_obj, const Tensor& grad_conf_noobj, const Tensor& grad_class, const int class_num, const float lambda_xy, const float lambda_wh, const float lambda_conf_obj, const float lambda_conf_noobj, const float lambda_class) { const int n = pred_x.dims()[0]; const int an_num = pred_x.dims()[1]; const int h = pred_x.dims()[2]; const int w = pred_x.dims()[3]; const int attr_num = class_num + 5; auto grad_t = EigenTensor::From(*grad).setConstant(0.0); auto pred_x_t = EigenTensor::From(pred_x); auto pred_y_t = EigenTensor::From(pred_y); auto pred_conf_t = EigenTensor::From(pred_conf); auto pred_class_t = EigenTensor::From(pred_class); auto grad_x_t = EigenTensor::From(grad_x); auto grad_y_t = EigenTensor::From(grad_y); auto grad_w_t = EigenTensor::From(grad_w); auto grad_h_t = EigenTensor::From(grad_h); auto grad_conf_obj_t = EigenTensor::From(grad_conf_obj); auto grad_conf_noobj_t = EigenTensor::From(grad_conf_noobj); auto grad_class_t = EigenTensor::From(grad_class); for (int i = 0; i < n; i++) { for (int j = 0; j < an_num; j++) { for (int k = 0; k < h; k++) { for (int l = 0; l < w; l++) { grad_t(i, j * attr_num, k, l) = grad_x_t(i, j, k, l) * pred_x_t(i, j, k, l) * (1.0 - pred_x_t(i, j, k, l)) * loss * lambda_xy; grad_t(i, j * attr_num + 1, k, l) = grad_y_t(i, j, k, l) * pred_y_t(i, j, k, l) * (1.0 - pred_y_t(i, j, k, l)) * loss * lambda_xy; grad_t(i, j * attr_num + 2, k, l) = grad_w_t(i, j, k, l) * loss * lambda_wh; grad_t(i, j * attr_num + 3, k, l) = grad_h_t(i, j, k, l) * loss * lambda_wh; grad_t(i, j * attr_num + 4, k, l) = grad_conf_obj_t(i, j, k, l) * pred_conf_t(i, j, k, l) * (1.0 - pred_conf_t(i, j, k, l)) * loss * lambda_conf_obj; grad_t(i, j * attr_num + 4, k, l) += grad_conf_noobj_t(i, j, k, l) * pred_conf_t(i, j, k, l) * (1.0 - pred_conf_t(i, j, k, l)) * loss * lambda_conf_noobj; for (int c = 0; c < class_num; c++) { grad_t(i, j * attr_num + 5 + c, k, l) = grad_class_t(i, j, k, l, c) * pred_class_t(i, j, k, l, c) * (1.0 - pred_class_t(i, j, k, l, c)) * loss * lambda_class; } } } } } } template class Yolov3LossKernel : public framework::OpKernel { public: void Compute(const framework::ExecutionContext& ctx) const override { auto* input = ctx.Input("X"); auto* gt_boxes = ctx.Input("GTBox"); auto* loss = ctx.Output("Loss"); auto anchors = ctx.Attr>("anchors"); int class_num = ctx.Attr("class_num"); float ignore_thresh = ctx.Attr("ignore_thresh"); float lambda_xy = ctx.Attr("lambda_xy"); float lambda_wh = ctx.Attr("lambda_wh"); float lambda_conf_obj = ctx.Attr("lambda_conf_obj"); float lambda_conf_noobj = ctx.Attr("lambda_conf_noobj"); float lambda_class = ctx.Attr("lambda_class"); const int n = input->dims()[0]; const int h = input->dims()[2]; const int w = input->dims()[3]; const int an_num = anchors.size() / 2; Tensor pred_x, pred_y, pred_w, pred_h; Tensor pred_conf, pred_class; pred_x.mutable_data({n, an_num, h, w}, ctx.GetPlace()); pred_y.mutable_data({n, an_num, h, w}, ctx.GetPlace()); pred_w.mutable_data({n, an_num, h, w}, ctx.GetPlace()); pred_h.mutable_data({n, an_num, h, w}, ctx.GetPlace()); pred_conf.mutable_data({n, an_num, h, w}, ctx.GetPlace()); pred_class.mutable_data({n, an_num, h, w, class_num}, ctx.GetPlace()); CalcPredResult(*input, &pred_conf, &pred_class, &pred_x, &pred_y, &pred_w, &pred_h, an_num, class_num); Tensor obj_mask, noobj_mask; Tensor tx, ty, tw, th, tconf, tclass; obj_mask.mutable_data({n, an_num, h, w}, ctx.GetPlace()); noobj_mask.mutable_data({n, an_num, h, w}, ctx.GetPlace()); tx.mutable_data({n, an_num, h, w}, ctx.GetPlace()); ty.mutable_data({n, an_num, h, w}, ctx.GetPlace()); tw.mutable_data({n, an_num, h, w}, ctx.GetPlace()); th.mutable_data({n, an_num, h, w}, ctx.GetPlace()); tconf.mutable_data({n, an_num, h, w}, ctx.GetPlace()); tclass.mutable_data({n, an_num, h, w, class_num}, ctx.GetPlace()); PreProcessGTBox(*gt_boxes, ignore_thresh, anchors, h, &obj_mask, &noobj_mask, &tx, &ty, &tw, &th, &tconf, &tclass); Tensor obj_mask_expand; obj_mask_expand.mutable_data({n, an_num, h, w, class_num}, ctx.GetPlace()); ExpandObjMaskByClassNum(&obj_mask_expand, obj_mask); T loss_x = CalcMSEWithMask(pred_x, tx, obj_mask); T loss_y = CalcMSEWithMask(pred_y, ty, obj_mask); T loss_w = CalcMSEWithMask(pred_w, tw, obj_mask); T loss_h = CalcMSEWithMask(pred_h, th, obj_mask); T loss_conf_obj = CalcBCEWithMask(pred_conf, tconf, obj_mask); T loss_conf_noobj = CalcBCEWithMask(pred_conf, tconf, noobj_mask); T loss_class = CalcBCEWithMask(pred_class, tclass, obj_mask_expand); auto* loss_data = loss->mutable_data({1}, ctx.GetPlace()); loss_data[0] = lambda_xy * (loss_x + loss_y) + lambda_wh * (loss_w + loss_h) + lambda_conf_obj * loss_conf_obj + lambda_conf_noobj * loss_conf_noobj + lambda_class * loss_class; } }; template class Yolov3LossGradKernel : public framework::OpKernel { public: void Compute(const framework::ExecutionContext& ctx) const override { auto* input = ctx.Input("X"); auto* gt_boxes = ctx.Input("GTBox"); auto anchors = ctx.Attr>("anchors"); int class_num = ctx.Attr("class_num"); float ignore_thresh = ctx.Attr("ignore_thresh"); auto* input_grad = ctx.Output(framework::GradVarName("X")); auto* output_grad = ctx.Input(framework::GradVarName("Loss")); const T loss = output_grad->data()[0]; float lambda_xy = ctx.Attr("lambda_xy"); float lambda_wh = ctx.Attr("lambda_wh"); float lambda_conf_obj = ctx.Attr("lambda_conf_obj"); float lambda_conf_noobj = ctx.Attr("lambda_conf_noobj"); float lambda_class = ctx.Attr("lambda_class"); const int n = input->dims()[0]; const int c = input->dims()[1]; const int h = input->dims()[2]; const int w = input->dims()[3]; const int an_num = anchors.size() / 2; Tensor pred_x, pred_y, pred_w, pred_h; Tensor pred_conf, pred_class; pred_x.mutable_data({n, an_num, h, w}, ctx.GetPlace()); pred_y.mutable_data({n, an_num, h, w}, ctx.GetPlace()); pred_w.mutable_data({n, an_num, h, w}, ctx.GetPlace()); pred_h.mutable_data({n, an_num, h, w}, ctx.GetPlace()); pred_conf.mutable_data({n, an_num, h, w}, ctx.GetPlace()); pred_class.mutable_data({n, an_num, h, w, class_num}, ctx.GetPlace()); CalcPredResult(*input, &pred_conf, &pred_class, &pred_x, &pred_y, &pred_w, &pred_h, an_num, class_num); Tensor obj_mask, noobj_mask; Tensor tx, ty, tw, th, tconf, tclass; obj_mask.mutable_data({n, an_num, h, w}, ctx.GetPlace()); noobj_mask.mutable_data({n, an_num, h, w}, ctx.GetPlace()); tx.mutable_data({n, an_num, h, w}, ctx.GetPlace()); ty.mutable_data({n, an_num, h, w}, ctx.GetPlace()); tw.mutable_data({n, an_num, h, w}, ctx.GetPlace()); th.mutable_data({n, an_num, h, w}, ctx.GetPlace()); tconf.mutable_data({n, an_num, h, w}, ctx.GetPlace()); tclass.mutable_data({n, an_num, h, w, class_num}, ctx.GetPlace()); PreProcessGTBox(*gt_boxes, ignore_thresh, anchors, h, &obj_mask, &noobj_mask, &tx, &ty, &tw, &th, &tconf, &tclass); Tensor obj_mask_expand; obj_mask_expand.mutable_data({n, an_num, h, w, class_num}, ctx.GetPlace()); ExpandObjMaskByClassNum(&obj_mask_expand, obj_mask); Tensor grad_x, grad_y, grad_w, grad_h; Tensor grad_conf_obj, grad_conf_noobj, grad_class; grad_x.mutable_data({n, an_num, h, w}, ctx.GetPlace()); grad_y.mutable_data({n, an_num, h, w}, ctx.GetPlace()); grad_w.mutable_data({n, an_num, h, w}, ctx.GetPlace()); grad_h.mutable_data({n, an_num, h, w}, ctx.GetPlace()); grad_conf_obj.mutable_data({n, an_num, h, w}, ctx.GetPlace()); grad_conf_noobj.mutable_data({n, an_num, h, w}, ctx.GetPlace()); grad_class.mutable_data({n, an_num, h, w, class_num}, ctx.GetPlace()); T obj_mf = CalcMaskPointNum(obj_mask); T noobj_mf = CalcMaskPointNum(noobj_mask); T obj_expand_mf = CalcMaskPointNum(obj_mask_expand); CalcMSEGradWithMask(&grad_x, pred_x, tx, obj_mask, obj_mf); CalcMSEGradWithMask(&grad_y, pred_y, ty, obj_mask, obj_mf); CalcMSEGradWithMask(&grad_w, pred_w, tw, obj_mask, obj_mf); CalcMSEGradWithMask(&grad_h, pred_h, th, obj_mask, obj_mf); CalcBCEGradWithMask(&grad_conf_obj, pred_conf, tconf, obj_mask, obj_mf); CalcBCEGradWithMask(&grad_conf_noobj, pred_conf, tconf, noobj_mask, noobj_mf); CalcBCEGradWithMask(&grad_class, pred_class, tclass, obj_mask_expand, obj_expand_mf); input_grad->mutable_data({n, c, h, w}, ctx.GetPlace()); AddAllGradToInputGrad( input_grad, loss, pred_x, pred_y, pred_conf, pred_class, grad_x, grad_y, grad_w, grad_h, grad_conf_obj, grad_conf_noobj, grad_class, class_num, lambda_xy, lambda_wh, lambda_conf_obj, lambda_conf_noobj, lambda_class); } }; } // namespace operators } // namespace paddle