Caffe YoloV3 training patch (#7)

* Fix type issue for SSD inference detection file(.json) dump

* Add support for Yolov2 training(COCO Dataset)

* Add multi_label feature in data layer to support YoloV3

* Remove unused function

* Add YoloV3 loss layer

* Fix segmentation fault error while doing make runtest

Author: SowmyaDhanapal

data/coco/labelmap_yolo_coco.prototxt

@@ -41,12 +41,15 @@ class BaseDataLayer : public Layer<Dtype> {
   TransformationParameter transform_param_;
   shared_ptr<DataTransformer<Dtype> > data_transformer_;
   bool output_labels_;
+  bool box_label_ = false;
+  int top_size_;
 };
 
 template <typename Dtype>
 class Batch {
  public:
   Blob<Dtype> data_, label_, dim_;
+  vector<shared_ptr<Blob<Dtype> > > multi_label_;
 };
 
 template <typename Dtype>
@@ -64,7 +67,6 @@ class BasePrefetchingDataLayer :
       const vector<Blob<Dtype>*>& top);
   virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
       const vector<Blob<Dtype>*>& top);
-
  protected:
   virtual void InternalThreadEntry();
   virtual void load_batch(Batch<Dtype>* batch) = 0;

include/caffe/layers/base_data_layer.hpp

@@ -15,8 +15,6 @@ namespace caffe {
     Dtype Overlap(Dtype x1, Dtype w1, Dtype x2, Dtype w2);
     template <typename Dtype>
     Dtype Calc_iou(const std::vector<Dtype>& box, const std::vector<Dtype>& truth);
-    template <typename Dtype>
-    Dtype Calc_rmse(const std::vector<Dtype>& box, const std::vector<Dtype>& truth);
 
     template<typename Dtype>
     class YoloV2LossLayer: public LossLayer<Dtype> {

include/caffe/layers/yolo_v2_loss_layer.hpp

@@ -0,0 +1,55 @@
+#ifndef YOLOV3LOSSLAYER_H
+#define YOLOV3LOSSLAYER_H
+
+#include <vector>
+#include <google/protobuf/repeated_field.h>
+
+#include "caffe/blob.hpp"
+#include "caffe/layer.hpp"
+#include "caffe/proto/caffe.pb.h"
+
+#include "caffe/layers/loss_layer.hpp"
+
+namespace caffe {
+    template <typename Dtype>
+    Dtype Overlap(Dtype x1, Dtype w1, Dtype x2, Dtype w2);
+    template <typename Dtype>
+    Dtype Calc_iou(const std::vector<Dtype>& box, const std::vector<Dtype>& truth);
+
+    template<typename Dtype>
+    class YoloV3LossLayer: public LossLayer<Dtype> {
+    public:
+        explicit YoloV3LossLayer(const LayerParameter& param)
+            : LossLayer<Dtype>(param), diff_() {}
+
+        virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
+              const vector<Blob<Dtype>*>& top);
+
+        virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
+              const vector<Blob<Dtype>*>& top);
+
+        virtual inline const char* type() const { return "YoloV3Loss"; }
+
+    protected:
+        virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
+              const vector<Blob<Dtype>*>& top);
+        virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
+              const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);
+
+        vector<int> biases_;
+        vector<int> masks_;
+        int side_;
+        int num_classes_;
+        int num_boxes_;
+        int total_num_boxes_;
+        float ignore_thresh_;
+        float truth_thresh_;
+        int seen;
+        int net_w_;
+        int net_h_;
+
+        Blob<Dtype> diff_;
+    };
+}
+
+#endif // YOLOV3LOSSLAYER_H

include/caffe/layers/yolo_v3_loss_layer.hpp

@@ -230,7 +230,7 @@ bool ReadXMLToAnnotatedDatum(const string& labelname, const int img_height,
 
 bool ReadJSONToAnnotatedDatum(const string& labelname, const int img_height,
     const int img_width, const std::map<string, int>& name_to_label,
-    AnnotatedDatum* anno_datum);
+    AnnotatedDatum* anno_datum, const bool caffe_yolo = false);
 
 bool ReadTxtToAnnotatedDatum(const string& labelname, const int height,
     const int width, AnnotatedDatum* anno_datum);

include/caffe/util/io.hpp

@@ -66,7 +66,6 @@ void AnnotatedDataLayer<Dtype>::DataLayerSetUp(
   for (int i = 0; i < this->prefetch_.size(); ++i) {
     this->prefetch_[i]->data_.Reshape(top_shape);
   }
-  //LOG(INFO) << "output data size: " << top[0]->num() << ","
   LOG_IF(INFO, Caffe::root_solver())
       << "output data size: " << top[0]->num() << ","
       << top[0]->channels() << "," << top[0]->height() << ","
@@ -75,15 +74,13 @@ void AnnotatedDataLayer<Dtype>::DataLayerSetUp(
   if (this->output_labels_) {
     has_anno_type_ = anno_datum.has_type() || anno_data_param.has_anno_type();
     if(transform_param.has_caffe_yolo()) {
+      this->box_label_ = true;
       vector<int> label_shape(1, batch_size);
       if (param.side_size() > 0) {
         for (int i = 0; i < param.side_size(); ++i) {
           sides_.push_back(param.side(i));
         }
       }
-      if (sides_.size() == 0) {
-        sides_.push_back(7);
-      }
       CHECK_EQ(sides_.size(), top.size() - 1) << "side num not equal to top size";
       if (has_anno_type_) {
         anno_type_ = anno_datum.type();
@@ -93,23 +90,28 @@ void AnnotatedDataLayer<Dtype>::DataLayerSetUp(
           LOG(WARNING) << "type stored in AnnotatedDatum is shadowed.";
           anno_type_ = anno_data_param.anno_type();
         }
-        if (anno_type_ == AnnotatedDatum_AnnotationType_BBOX) { 
+        for (int i = 0; i < this->prefetch_.size(); ++i) {
+          this->prefetch_[i]->multi_label_.clear();
+        }
+        if (anno_type_ == AnnotatedDatum_AnnotationType_BBOX) {
           // Yolo label format
           for (int i = 0; i < sides_.size(); ++i) {
             vector<int> label_shape(1, batch_size);
             int label_size = sides_[i] * sides_[i] * (1 + 1 + 1 + 4);
             label_shape.push_back(label_size);
             top[i+1]->Reshape(label_shape);
+            for (int j = 0; j < this->prefetch_.size(); ++j) {
+              shared_ptr<Blob<Dtype> > tmp_blob;
+              tmp_blob.reset(new Blob<Dtype>(label_shape));
+              this->prefetch_[j]->multi_label_.push_back(tmp_blob);
+            }
           }
         } else {
           LOG(FATAL) << "Unknown annotation type.";
         }
       } else {
         label_shape[0] = batch_size;
       }
-      for (int i = 0; i < this->prefetch_.size(); ++i) { 
-        this->prefetch_[i]->label_.Reshape(label_shape);
-      }
     }
     else {
       vector<int> label_shape(4, 1);
@@ -214,7 +216,10 @@ void AnnotatedDataLayer<Dtype>::load_batch(Batch<Dtype>* batch) {
   // Store transformed annotation.
   map<int, vector<AnnotationGroup> > all_anno;
   int num_bboxes = 0;
-
+  vector<Dtype*> top_label;
+  for (int i = 0; i < sides_.size(); ++i) {
+    top_label.push_back(batch->multi_label_[i]->mutable_cpu_data());
+  }
   for (int item_id = 0; item_id < batch_size; ++item_id) {
     timer.Start();
     // get a anno_datum
@@ -333,47 +338,47 @@ void AnnotatedDataLayer<Dtype>::load_batch(Batch<Dtype>* batch) {
           for (int i = 0; i < sides_.size(); ++i) {
             side = sides_[i];
             count = sides_[i] * sides_[i] * (1 + 1 + 1 + 4);
-          }
-          label_shape[0] = batch_size;
-          label_shape[1] = count;
-          batch->label_.Reshape(label_shape);
-          Dtype* top_label = batch->label_.mutable_cpu_data();
-          const vector<AnnotationGroup>& anno_vec = all_anno[item_id];
-          label_offset = count * item_id;
-          top_label = top_label + label_offset;
-          int locations = pow(side, 2);
-          CHECK_EQ(count, locations * 7) << "side and count not match";
-          // difficult
-          caffe_set(locations, Dtype(0), top_label);
-          // isobj
-          caffe_set(locations, Dtype(0), top_label + locations);
-          // class label
-          caffe_set(locations, Dtype(-1), top_label + locations * 2);
-          // bounding box
-          caffe_set(locations*4, Dtype(0), top_label + locations * 3);
-          for (int g = 0; g < anno_vec.size(); ++g) {
-            const AnnotationGroup& anno_group = anno_vec[g];
-            for (int a = 0; a < anno_group.annotation_size(); ++a) {
-              const Annotation& anno = anno_group.annotation(a);
-              const NormalizedBBox& bbox = anno.bbox();
-              float class_label = anno_group.group_label();
-              float x = bbox.x_center();
-              float y = bbox.y_center();
-              int x_index = floor(x * side);
-              int y_index = floor(y * side);
-              x_index = std::min(x_index, side - 1);
-              y_index = std::min(y_index, side - 1);
-              int dif_index = side * y_index + x_index;
-              int obj_index = locations + dif_index;
-              int class_index = locations * 2 + dif_index;
-              int cor_index = locations * 3 + dif_index * 4;
-              top_label[dif_index] = bbox.difficult();
-              top_label[obj_index] = 1;
-              top_label[class_index] = class_label;
-              top_label[cor_index + 0] = bbox.x_center();
-              top_label[cor_index + 1] = bbox.y_center();
-              top_label[cor_index + 2] = bbox.width();
-              top_label[cor_index + 3] = bbox.height();
+            label_shape[0] = batch_size;
+            label_shape[1] = count;
+            batch->multi_label_[i]->Reshape(label_shape);
+            top_label[i] = batch->multi_label_[i]->mutable_cpu_data();
+            const vector<AnnotationGroup>& anno_vec = all_anno[item_id];
+            label_offset = count * item_id;
+            top_label[i] = top_label[i] + label_offset;
+            int locations = pow(side, 2);
+            CHECK_EQ(count, locations * 7) << "side and count not match";
+            // difficult
+            caffe_set(locations, Dtype(0), top_label[i]);
+            // isobj
+            caffe_set(locations, Dtype(0), top_label[i] + locations);
+            // class label
+            caffe_set(locations, Dtype(-1), top_label[i] + locations * 2);
+            // bounding box
+            caffe_set(locations*4, Dtype(0), top_label[i] + locations * 3);
+            for (int g = 0; g < anno_vec.size(); ++g) {
+              const AnnotationGroup& anno_group = anno_vec[g];
+              for (int a = 0; a < anno_group.annotation_size(); ++a) {
+                const Annotation& anno = anno_group.annotation(a);
+                const NormalizedBBox& bbox = anno.bbox();
+                float class_label = anno_group.group_label();
+                float x = bbox.x_center();
+                float y = bbox.y_center();
+                int x_index = floor(x * side);
+                int y_index = floor(y * side);
+                x_index = std::min(x_index, side - 1);
+                y_index = std::min(y_index, side - 1);
+                int dif_index = side * y_index + x_index;
+                int obj_index = locations + dif_index;
+                int class_index = locations * 2 + dif_index;
+                int cor_index = locations * 3 + dif_index * 4;
+                top_label[i][dif_index] = bbox.difficult();
+                top_label[i][obj_index] = 1;
+                top_label[i][class_index] = class_label;
+                top_label[i][cor_index + 0] = bbox.x_center();
+                top_label[i][cor_index + 1] = bbox.y_center();
+                top_label[i][cor_index + 2] = bbox.width();
+                top_label[i][cor_index + 3] = bbox.height();
+              }
             }
           }
         }
@@ -401,7 +406,6 @@ void AnnotatedDataLayer<Dtype>::load_batch(Batch<Dtype>* batch) {
     //reader_.free().push(const_cast<AnnotatedDatum*>(&anno_datum));
     Next();
   }
-
   // Store "rich" annotation if needed.
   if (this->output_labels_ && has_anno_type_) {
     if(anno_type_ == AnnotatedDatum_AnnotationType_BBOX) {

src/caffe/layers/annotated_data_layer.cpp

@@ -56,15 +56,30 @@ void BasePrefetchingDataLayer<Dtype>::LayerSetUp(
   for (int i = 0; i < prefetch_.size(); ++i) {
     prefetch_[i]->data_.mutable_cpu_data();
     if (this->output_labels_) {
-      prefetch_[i]->label_.mutable_cpu_data();
+      if (this->box_label_) {
+        this->top_size_ = top.size();
+        for (int j = 0; j < top.size() - 1; ++j) {
+          prefetch_[i]->multi_label_[j]->mutable_cpu_data();
+        }
+      }
+      else {
+        prefetch_[i]->label_.mutable_cpu_data();
+      }
     }
   }
 #ifndef CPU_ONLY
   if (Caffe::mode() == Caffe::GPU) {
     for (int i = 0; i < prefetch_.size(); ++i) {
       prefetch_[i]->data_.mutable_gpu_data();
       if (this->output_labels_) {
-        prefetch_[i]->label_.mutable_gpu_data();
+        if (this->box_label_) {
+          for (int j = 0; j < top.size() - 1; ++j) {
+            prefetch_[i]->multi_label_[j]->mutable_gpu_data();
+          }
+        }
+        else {
+          prefetch_[i]->label_.mutable_gpu_data();
+        }
       }
     }
   }
@@ -92,7 +107,14 @@ void BasePrefetchingDataLayer<Dtype>::InternalThreadEntry() {
       if (Caffe::mode() == Caffe::GPU) {
         batch->data_.data().get()->async_gpu_push(stream);
         if (this->output_labels_) {
-          batch->label_.data().get()->async_gpu_push(stream);
+          if(this->box_label_) {
+            for (int j = 0; j < this->top_size_ - 1; ++j) {
+              batch->multi_label_[j]->data().get()->async_gpu_push(stream);
+            }
+          }
+          else {
+            batch->label_.data().get()->async_gpu_push(stream);
+          }
         }
         CUDA_CHECK(cudaStreamSynchronize(stream));
       }
@@ -120,9 +142,17 @@ void BasePrefetchingDataLayer<Dtype>::Forward_cpu(
   top[0]->ReshapeLike(prefetch_current_->data_);
   top[0]->set_cpu_data(prefetch_current_->data_.mutable_cpu_data());
   if (this->output_labels_) {
-    // Reshape to loaded labels.
-    top[1]->ReshapeLike(prefetch_current_->label_);
-    top[1]->set_cpu_data(prefetch_current_->label_.mutable_cpu_data());
+    if(this->box_label_) {
+      for (int j = 0; j < top.size() - 1; ++j) {
+        top[j+1]->ReshapeLike(*(prefetch_current_->multi_label_[j]));
+        top[j+1]->set_cpu_data(prefetch_current_->multi_label_[j]->mutable_cpu_data());
+      }
+    }
+    else{
+      // Reshape to loaded labels.
+      top[1]->ReshapeLike(prefetch_current_->label_);
+      top[1]->set_cpu_data(prefetch_current_->label_.mutable_cpu_data());
+    }
   }
 }
 

src/caffe/layers/base_data_layer.cpp

@@ -15,9 +15,17 @@ void BasePrefetchingDataLayer<Dtype>::Forward_gpu(
   top[0]->ReshapeLike(prefetch_current_->data_);
   top[0]->set_gpu_data(prefetch_current_->data_.mutable_gpu_data());
   if (this->output_labels_) {
-    // Reshape to loaded labels.
-    top[1]->ReshapeLike(prefetch_current_->label_);
-    top[1]->set_gpu_data(prefetch_current_->label_.mutable_gpu_data());
+    if(this->box_label_) {
+      for (int j = 0; j < top.size() - 1; ++j) {
+        top[j+1]->ReshapeLike(*(prefetch_current_->multi_label_[j]));
+        top[j+1]->set_gpu_data(prefetch_current_->multi_label_[j]->mutable_gpu_data());
+      }
+    }
+    else {
+      // Reshape to loaded labels.
+      top[1]->ReshapeLike(prefetch_current_->label_);
+      top[1]->set_gpu_data(prefetch_current_->label_.mutable_gpu_data());
+    }
   }
 }
 

src/caffe/layers/base_data_layer.cu

@@ -68,15 +68,6 @@ namespace caffe {
         return x;
     }
 
-    template <typename Dtype>
-    Dtype Calc_rmse(const Box<Dtype>& truth, const Box<Dtype>& box, Dtype &coord_loss, Dtype &area_loss, float scale) {
-        float coord_ = scale * (abs(box.x-truth.x) + abs(box.y-truth.y));
-        float area_  = scale * (abs(box.w-truth.w) + abs(box.h-truth.h));
-        coord_loss += coord_;
-        area_loss  += area_;
-        return (coord_ + area_);
-    }
-
    template <typename Dtype>
     float delta_region_box(Box<Dtype>& truth_box, Dtype* input, 
             std::vector<float>& biases, int n, int side, Dtype* diff,