ExternalSource operator fix (#106)

* Fix ExternalSource op and add an example Signed-off-by: Serge Panev <spanev@nvidia.com> * Thread safety for ExernalSource operator Signed-off-by: Serge Panev <spanev@nvidia.com> * Set ExternalSource dtor to default Signed-off-by: Serge Panev <spanev@nvidia.com> * Switch from std::atomic to std::mutex in ExternalSource Signed-off-by: Serge Panev <spanev@nvidia.com> * Extend esternal input example Signed-off-by: Janusz Lisiecki <jlisiecki@nvidia.com> * Updated ExternalSource docs Signed-off-by: Janusz Lisiecki <jlisiecki@nvidia.com>
NVIDIA · Aug 31, 2018 · 58ab4a9 · 58ab4a9
1 parent 6c45ef3
commit 58ab4a9
Show file tree

Hide file tree

Showing 6 changed files with 280 additions and 9 deletions.
diff --git a/dali/pipeline/data/tensor.h b/dali/pipeline/data/tensor.h
@@ -60,6 +60,19 @@ class Tensor : public Buffer<Backend> {
         other.raw_data(), this->size(), stream);
   }
 
+  /**
+   * @brief Loads the Tensor at index idx from the input TensorList.
+   */
+  template <typename InBackend>
+  inline void Copy(const TensorList<InBackend> &other, int idx, cudaStream_t stream) {
+    shape_ = other.tensor_shape(idx);
+    device_ = other.device_id();
+    this->set_type(other.type());
+    this->Resize(shape_);
+    type_.template Copy<Backend, InBackend>(this->raw_mutable_data(),
+        other.raw_tensor(idx), this->size(), stream);
+  }
+
   template <typename InBackend>
   inline void ResizeLike(const Tensor<InBackend> &other) {
     Resize(other.shape());

diff --git a/dali/pipeline/operators/util/external_source.cc b/dali/pipeline/operators/util/external_source.cc
@@ -20,19 +20,30 @@ template<>
 void ExternalSource<CPUBackend>::RunImpl(SampleWorkspace *ws, const int idx) {
   // Wrap the output tensor around our data
   auto output = ws->Output<CPUBackend>(idx);
+  cudaStream_t stream = ws->has_stream() ? ws->stream() : 0;
   if (data_in_tl_) {
-    output->ShareData(&tl_data_, ws->data_idx());
+    output->Copy(tl_data_, ws->data_idx(), stream);
   } else {
     DALI_ENFORCE_VALID_INDEX(ws->data_idx(), t_data_.size());
     auto &data = t_data_[ws->data_idx()];
-    output->ShareData(&data);
+    output->Copy(data, stream);
+  }
+
+  std::unique_lock<std::mutex> l(samples_processed_m_);
+  if (++samples_processed_ >= batch_size_) {
+    samples_processed_ = 0;
+    busy_ = false;
+    cv_.notify_one();
   }
 }
 
 DALI_REGISTER_OPERATOR(ExternalSource, ExternalSource<CPUBackend>, CPU);
 
 DALI_SCHEMA(ExternalSource)
-  .DocStr(R"code(Allows externally provided data to be passed as an input to the pipeline)code")
+  .DocStr(R"code(Allows externally provided data to be passed as an input to the pipeline,
+           see :meth:`nvidia.dali.pipeline.Pipeline.feed_input` and
+           :meth:`nvidia.dali.pipeline.Pipeline.iter_setup`. Currenlty this operator is not
+           supported in TensorFlow.)code")
   .NumInput(0)
   .NumOutput(1);
 

diff --git a/dali/pipeline/operators/util/external_source.cu b/dali/pipeline/operators/util/external_source.cu
@@ -19,8 +19,11 @@ namespace dali {
 template<>
 void ExternalSource<GPUBackend>::RunImpl(DeviceWorkspace *ws, const int idx) {
   DALI_ENFORCE(data_in_tl_, "Cannot feed non-contiguous data to GPU op.");
+
   auto output = ws->Output<GPUBackend>(idx);
-  output->ShareData(&tl_data_);
+  output->Copy(tl_data_, (ws->has_stream() ? ws->stream() : 0));
+  busy_ = false;
+  cv_.notify_all();
 }
 
 DALI_REGISTER_OPERATOR(ExternalSource, ExternalSource<GPUBackend>, GPU);

diff --git a/dali/pipeline/operators/util/external_source.h b/dali/pipeline/operators/util/external_source.h
@@ -15,8 +15,11 @@
 #ifndef DALI_PIPELINE_OPERATORS_UTIL_EXTERNAL_SOURCE_H_
 #define DALI_PIPELINE_OPERATORS_UTIL_EXTERNAL_SOURCE_H_
 
+#include <atomic>
 #include <string>
 #include <vector>
+#include <condition_variable>
+#include <mutex>
 
 #include "dali/pipeline/operators/operator.h"
 
@@ -31,11 +34,13 @@ template <typename Backend>
 class ExternalSource : public Operator<Backend> {
  public:
   inline explicit ExternalSource(const OpSpec &spec) :
-    Operator<Backend>(spec) {
+    Operator<Backend>(spec),
+    samples_processed_(0),
+    busy_(false) {
     output_name_ = spec.Output(0);
   }
 
-  virtual inline ~ExternalSource() = default;
+  inline ~ExternalSource() = default;
 
   inline string name() const override {
     return "ExternalSource (" + output_name_ + ")";
@@ -49,8 +54,11 @@ class ExternalSource : public Operator<Backend> {
     // Note: If we create a GPU source, we will need to figure
     // out what stream we want to do this copy in. CPU we can
     // pass anything as it is ignored.
+    std::unique_lock<std::mutex> lock(m_);
+    cv_.wait(lock,  [this]{return !this->busy_;});
     tl_data_.Copy(tl, 0);
     data_in_tl_ = true;
+    busy_ = true;
   }
 
   /**
@@ -61,11 +69,16 @@ class ExternalSource : public Operator<Backend> {
     // Note: If we create a GPU source, we will need to figure
     // out what stream we want to do this copy in. CPU we can
     // pass anything as it is ignored.
+
+    std::unique_lock<std::mutex> lock(m_);
+    cv_.wait(lock,  [this]{return !this->busy_;});
+
     t_data_.resize(t.size());
     for (size_t i = 0; i < t.size(); ++i) {
       t_data_[i].Copy(t[i], 0);
     }
     data_in_tl_ = false;
+    busy_ = true;
   }
 
   DISABLE_COPY_MOVE_ASSIGN(ExternalSource);
@@ -75,8 +88,15 @@ class ExternalSource : public Operator<Backend> {
 
   string output_name_;
   TensorList<Backend> tl_data_;
-  vector<Tensor<Backend>> t_data_;
-  bool data_in_tl_ = true;
+  std::vector<Tensor<Backend>> t_data_;
+  bool data_in_tl_;
+
+  std::atomic<int> samples_processed_;
+
+  bool busy_;
+  std::condition_variable cv_;
+  std::mutex m_;
+  std::mutex samples_processed_m_;
 };
 
 }  // namespace dali

diff --git a/docs/examples/external_input.ipynb b/docs/examples/external_input.ipynb
diff --git a/docs/examples/images/file_list.txt b/docs/examples/images/file_list.txt
@@ -19,4 +19,3 @@ kitten/cat_7.jpg 1
 kitten/cat_8.jpg 1
 kitten/cat_9.jpg 1
 kitten/cat_10.jpg 1
-