Changes to dali/util/numpy

- Make `numpy.h` available to users - Extend `ReadTensor` to also accept a `bool pinned` argument - See #5337 for details Signed-off-by: Francesco Versaci <francesco.versaci@gmail.com>
NVIDIA · Apr 5, 2024 · 1c973d0 · 1c973d0
1 parent 43f2671
commit 1c973d0
Show file tree

Hide file tree

Showing 5 changed files with 6 additions and 4 deletions.
diff --git a/dali/kernels/slice/slice_flip_normalize_gpu_test.cu b/dali/kernels/slice/slice_flip_normalize_gpu_test.cu
@@ -80,7 +80,7 @@ class SliceFlipNormalizeGPUTest : public ::testing::Test {
 
   void LoadTensor(Tensor<CPUBackend> &tensor, const std::string& path_npy) {
     auto stream = FileStream::Open(path_npy, false, false);
-    tensor = ::dali::numpy::ReadTensor(stream.get());
+    tensor = ::dali::numpy::ReadTensor(stream.get(), true);
   }
 
   template <typename T, int ndim>

diff --git a/dali/operators/imgcodec/decoder_test_helper.h b/dali/operators/imgcodec/decoder_test_helper.h
@@ -238,7 +238,7 @@ inline void AssertClose(const Tensor<CPUBackend> &img, const Tensor<CPUBackend>
 }
 
 inline Tensor<CPUBackend> ReadReference(InputStream *src, TensorLayout layout = "HWC") {
-  auto tensor = numpy::ReadTensor(src);
+  auto tensor = numpy::ReadTensor(src, true);
   tensor.SetLayout(layout);
   return tensor;
 }

diff --git a/dali/util/CMakeLists.txt b/dali/util/CMakeLists.txt
@@ -22,6 +22,7 @@ set(DALI_INST_HDRS ${DALI_INST_HDRS}
   "${CMAKE_CURRENT_SOURCE_DIR}/ocv.h"
   "${CMAKE_CURRENT_SOURCE_DIR}/random_crop_generator.h"
   "${CMAKE_CURRENT_SOURCE_DIR}/thread_safe_queue.h"
+  "${CMAKE_CURRENT_SOURCE_DIR}/numpy.h"
   "${CMAKE_CURRENT_SOURCE_DIR}/user_stream.h")
 
 set(DALI_SRCS ${DALI_SRCS}

diff --git a/dali/util/numpy.cc b/dali/util/numpy.cc
@@ -294,12 +294,13 @@ size_t HeaderData::nbytes() const {
   return type_info ? type_info->size() * size() : 0_uz;
 }
 
-Tensor<CPUBackend> ReadTensor(InputStream *src) {
+  Tensor<CPUBackend> ReadTensor(InputStream *src, bool pinned) {
   numpy::HeaderData header;
   numpy::ParseHeader(header, src);
   src->SeekRead(header.data_offset, SEEK_SET);
 
   Tensor<CPUBackend> data;
+  data.set_pinned(pinned);
   data.Resize(header.shape, header.type());
   auto ret = src->Read(static_cast<uint8_t*>(data.raw_mutable_data()), header.nbytes());
   DALI_ENFORCE(ret == header.nbytes(), "Failed to read numpy file");

diff --git a/dali/util/numpy.h b/dali/util/numpy.h
@@ -55,7 +55,7 @@ DLL_PUBLIC void FromFortranOrder(SampleView<CPUBackend> output, ConstSampleView<
 
 DLL_PUBLIC void ParseHeaderContents(HeaderData& target, const std::string &header);
 
-DLL_PUBLIC Tensor<CPUBackend> ReadTensor(InputStream *src);
+DLL_PUBLIC Tensor<CPUBackend> ReadTensor(InputStream *src, bool pinned);
 
 }  // namespace numpy
 }  // namespace dali