Tensor layout checking (NVIDIA#32)

* SupportWorkspace methods templated for consistency Signed-off-by: Serge Panev <[email protected]> * Add Layout attribute to Tensor/TL - layout checking in Operator Signed-off-by: Serge Panev <[email protected]> * Enforces layout in layout-dependant Operators Signed-off-by: Serge Panev <[email protected]>
zzm422 · Jul 9, 2018 · fdccf61 · fdccf61
1 parent 58f96e3
commit fdccf61
Show file tree

Hide file tree

Showing 12 changed files with 93 additions and 13 deletions.
diff --git a/dali/pipeline/data/tensor.h b/dali/pipeline/data/tensor.h
@@ -35,7 +35,7 @@ namespace dali {
 template <typename Backend>
 class Tensor : public Buffer<Backend> {
  public:
-  inline Tensor() {}
+  inline Tensor() : layout_(DALI_NHWC) {}
   inline ~Tensor() = default;
 
   /**
@@ -116,8 +116,8 @@ class Tensor : public Buffer<Backend> {
 
   /**
    * @brief Wraps the data owned by the input tensor. The input
-   * tensor must have a valid type. If sucessful, the tensor 
-   * object will wrap the target data and assume the datatype 
+   * tensor must have a valid type. If sucessful, the tensor
+   * object will wrap the target data and assume the datatype
    * and shape of the data stored in the Tensor.
    *
    * If the input does not store any data, shares_data_ is left
@@ -292,8 +292,17 @@ class Tensor : public Buffer<Backend> {
     return *this;
   }
 
+  inline DALITensorLayout GetLayout() const {
+    return layout_;
+  }
+
+  inline void SetLayout(DALITensorLayout layout) {
+    layout_ = layout;
+  }
+
  protected:
   vector<Index> shape_;
+  DALITensorLayout layout_;
 
   USE_BUFFER_MEMBERS();
 };

diff --git a/dali/pipeline/data/tensor_list.h b/dali/pipeline/data/tensor_list.h
@@ -40,7 +40,7 @@ typedef vector<Index> Dims;
 template <typename Backend>
 class DLL_PUBLIC TensorList : public Buffer<Backend> {
  public:
-  DLL_PUBLIC TensorList() {}
+  DLL_PUBLIC TensorList() : layout_(DALI_NHWC) {}
   DLL_PUBLIC ~TensorList() = default;
 
   /**
@@ -273,12 +273,21 @@ class DLL_PUBLIC TensorList : public Buffer<Backend> {
 
   DISABLE_COPY_MOVE_ASSIGN(TensorList);
 
+  inline DALITensorLayout GetLayout() const {
+    return layout_;
+  }
+
+  inline void SetLayout(DALITensorLayout layout) {
+    layout_ = layout;
+  }
+
  protected:
   // We store a set of dimension for each tensor in the list.
   // We also pre-compute the offsets of each tensor in the
   // underlying allocation for random access
   vector<Dims> shape_;
   vector<Index> offsets_;
+  DALITensorLayout layout_;
 
   USE_BUFFER_MEMBERS();
 };

diff --git a/dali/pipeline/operators/fused/crop_mirror_normalize.cu b/dali/pipeline/operators/fused/crop_mirror_normalize.cu
@@ -282,6 +282,9 @@ void CropMirrorNormalize<GPUBackend>::DataDependentSetup(DeviceWorkspace *ws, co
   // Resize the output data
   output->Resize(output_shape);
 
+  // Set the layout of the output data
+  output->SetLayout(output_layout_);
+
   // Copy strides to gpu
   input_strides_gpu_.Copy(input_strides_, ws->stream());
 

diff --git a/dali/pipeline/operators/fused/normalize_permute.cc b/dali/pipeline/operators/fused/normalize_permute.cc
@@ -33,6 +33,7 @@ namespace dali {
 
     // Output is CHW
     output->Resize({C_, H_, W_});
+    output->SetLayout(DALI_NCHW);
     if (output_type_ == DALI_FLOAT) {
       CPURunHelper<float>(input, output);
     } else {
@@ -88,6 +89,7 @@ DALI_SCHEMA(NormalizePermute)
       Mean pixel values for image normalization)code")
   .AddArg("std",
       R"code(`list of float`
-      Standard deviation values for image normalization)code");
+      Standard deviation values for image normalization)code")
+  .EnforceInputLayout(DALI_NHWC);
 
 }  // namespace dali
diff --git a/dali/pipeline/operators/fused/resize_crop_mirror.cc b/dali/pipeline/operators/fused/resize_crop_mirror.cc
@@ -42,7 +42,8 @@ DALI_SCHEMA(ResizeCropMirror)
   .AddOptionalArg("crop_pos_y",
       "Vertical position of the crop in image coordinates (0.0 - 1.0)",
       0.5f)
-  .AddOptionalArg("mirror", "Mask for horizontal flip", 0);
+  .AddOptionalArg("mirror", "Mask for horizontal flip", 0)
+  .EnforceInputLayout(DALI_NHWC);
 
 DALI_REGISTER_OPERATOR(FastResizeCropMirror, FastResizeCropMirror<CPUBackend>, CPU);
 
@@ -71,6 +72,7 @@ DALI_SCHEMA(FastResizeCropMirror)
   .AddOptionalArg("crop_pos_y",
       "Vertical position of the crop in image coordinates (0.0 - 1.0)",
       0.5f)
-  .AddOptionalArg("mirror", "Mask for horizontal flip", 0);
+  .AddOptionalArg("mirror", "Mask for horizontal flip", 0)
+  .EnforceInputLayout(DALI_NHWC);
 
 }  // namespace dali
diff --git a/dali/pipeline/operators/op_schema.h b/dali/pipeline/operators/op_schema.h
@@ -36,7 +36,8 @@ class DLL_PUBLIC OpSchema {
 
   DLL_PUBLIC explicit inline OpSchema(const std::string &name)
     : name_(name),
-      allow_multiple_input_sets_(false) {
+      allow_multiple_input_sets_(false),
+      enforce_layout_(false) {
     // Fill internal arguments
     internal_arguments_["num_threads"] = std::make_pair("Number of CPU threads in a thread pool",
         Value::construct(-1));
@@ -146,6 +147,12 @@ class DLL_PUBLIC OpSchema {
     return *this;
   }
 
+  DLL_PUBLIC inline OpSchema& EnforceInputLayout(DALITensorLayout layout) {
+    layout_ = layout;
+    enforce_layout_ = true;
+    return *this;
+  }
+
   /**
    * @brief Adds an optional non-vector argument to op
    */
@@ -226,6 +233,14 @@ class DLL_PUBLIC OpSchema {
     return allow_multiple_input_sets_;
   }
 
+  DLL_PUBLIC inline bool EnforceInputLayout() const {
+    return enforce_layout_;
+  }
+
+  DLL_PUBLIC inline DALITensorLayout InputLayout() const {
+    return layout_;
+  }
+
   DLL_PUBLIC inline bool HasOutputFn() const {
     return static_cast<bool>(output_fn_);
   }
@@ -289,6 +304,9 @@ class DLL_PUBLIC OpSchema {
   bool allow_multiple_input_sets_;
   vector<string> parents_;
 
+  bool enforce_layout_;
+  DALITensorLayout layout_;
+
   std::map<std::string, std::string> arguments_;
   std::map<std::string, std::pair<std::string, Value*> > optional_arguments_;
   std::map<std::string, std::pair<std::string, Value*> > internal_arguments_;

diff --git a/dali/pipeline/operators/operator.h b/dali/pipeline/operators/operator.h
@@ -38,6 +38,37 @@ enum DALIOpType {
   DALI_SUPPORT = 3
 };
 
+template <typename InputType>
+inline void CheckInputLayout(const InputType& input, const OpSpec& spec) {
+  auto schema = SchemaRegistry::GetSchema(spec.name());
+  if (schema.EnforceInputLayout()) {
+    DALI_ENFORCE(input.GetLayout() == schema.InputLayout());
+  }
+}
+
+template <typename Workspace>
+inline void CheckInputLayouts(const Workspace *ws, const OpSpec &spec) {
+  for (int i = 0; i < spec.NumRegularInput(); ++i) {
+    auto& input = ws->template Input<CPUBackend>(i);
+    CheckInputLayout(input, spec);
+  }
+}
+
+template <>
+inline void CheckInputLayouts(const DeviceWorkspace *ws, const OpSpec &spec) {
+  for (int i = 0; i < spec.NumRegularInput(); ++i) {
+    if (ws->InputIsType<CPUBackend>(i)) {
+      auto& input = ws->Input<CPUBackend>(i);
+      CheckInputLayout(input, spec);
+    } else if (ws->InputIsType<GPUBackend>(i)) {
+      auto& input = ws->Input<GPUBackend>(i);
+      CheckInputLayout(input, spec);
+    } else {
+      DALI_FAIL("Input has an unkown backend");
+    }
+  }
+}
+
 /**
  * @brief Baseclass for the basic unit of computation in the pipeline.
  *
@@ -141,6 +172,7 @@ class Operator : public OperatorBase {
 
   using OperatorBase::Run;
   void Run(Workspace<Backend> *ws) override {
+    CheckInputLayouts(ws, spec_);
     SetupSharedSampleParams(ws);
 
     for (int i = 0; i < input_sets_; ++i) {

diff --git a/dali/pipeline/operators/resize/random_resized_crop.cc b/dali/pipeline/operators/resize/random_resized_crop.cc
@@ -45,7 +45,8 @@ DALI_SCHEMA(RandomResizedCrop)
       Size of resized image)code")
   .AddOptionalArg("num_attempts",
       R"code(`int`
-      Maximum number of attempts used to choose random area and aspect ratio)code", 10);
+      Maximum number of attempts used to choose random area and aspect ratio)code", 10)
+  .EnforceInputLayout(DALI_NHWC);
 
 template<>
 struct RandomResizedCrop<CPUBackend>::Params {

diff --git a/dali/pipeline/operators/support/random/coin_flip.cc b/dali/pipeline/operators/support/random/coin_flip.cc
@@ -19,7 +19,7 @@ namespace dali {
 
 void CoinFlip::RunImpl(SupportWorkspace * ws, const int idx) {
   DALI_ENFORCE(idx == 0, "CoinFlip does not support multiple input sets.");
-  auto *output = ws->Output(idx);
+  auto *output = ws->Output<CPUBackend>(idx);
   output->Resize({batch_size_});
 
   int * out_data = output->template mutable_data<int>();

diff --git a/dali/pipeline/operators/support/random/uniform.cc b/dali/pipeline/operators/support/random/uniform.cc
@@ -21,7 +21,7 @@ namespace dali {
 
 void Uniform::RunImpl(SupportWorkspace * ws, const int idx) {
   DALI_ENFORCE(idx == 0, "Uniform does not support multiple input sets.");
-  auto *output = ws->Output(idx);
+  auto *output = ws->Output<CPUBackend>(idx);
   output->Resize({batch_size_});
 
   float * out_data = output->template mutable_data<float>();

diff --git a/dali/pipeline/workspace/support_workspace.cc b/dali/pipeline/workspace/support_workspace.cc
@@ -16,6 +16,7 @@
 
 namespace dali {
 
+template <>
 const Tensor<CPUBackend>& SupportWorkspace::Input(int idx) const {
   DALI_ENFORCE_VALID_INDEX(idx, input_index_map_.size());
   auto tensor_meta = input_index_map_[idx];
@@ -24,6 +25,7 @@ const Tensor<CPUBackend>& SupportWorkspace::Input(int idx) const {
   return *cpu_inputs_[tensor_meta.second];
 }
 
+template <>
 Tensor<CPUBackend>* SupportWorkspace::Output(int idx) {
   DALI_ENFORCE_VALID_INDEX(idx, output_index_map_.size());
   auto tensor_meta = output_index_map_[idx];

diff --git a/dali/pipeline/workspace/support_workspace.h b/dali/pipeline/workspace/support_workspace.h
@@ -44,12 +44,14 @@ class DLL_PUBLIC SupportWorkspace : public WorkspaceBase<SupportInputType, Suppo
   /**
    * @brief Returns the input Tensor at index `idx`.
    */
-  DLL_PUBLIC const Tensor<CPUBackend>& Input(int idx) const;
+  template <typename Backend>
+  DLL_PUBLIC const Tensor<Backend>& Input(int idx) const;
 
   /**
    * @brief Returns the output Tensor at index `idx`.
    */
-  DLL_PUBLIC Tensor<CPUBackend>* Output(int idx);
+  template <typename Backend>
+  DLL_PUBLIC Tensor<Backend>* Output(int idx);
 };
 
 }  // namespace dali