Move models out of data and install them with onnxrt (#369)

astro-informatics · Nov 14, 2024 · 99a63e0 · 99a63e0
1 parent e384883
commit 99a63e0
Show file tree

Hide file tree

Showing 6 changed files with 7 additions and 3 deletions.
diff --git a/cmake_files/dependencies.cmake b/cmake_files/dependencies.cmake
@@ -28,6 +28,7 @@ find_package(yaml-cpp REQUIRED)
 
 if (onnxrt)
   find_package(onnxruntime REQUIRED)
+  install(DIRECTORY ${CMAKE_SOURCE_DIR}/models DESTINATION .)
 endif()
 
 find_package(sopt REQUIRED)

diff --git a/cpp/tests/algo_factory.cc b/cpp/tests/algo_factory.cc
@@ -228,7 +228,7 @@ TEST_CASE("tf_fb_factory") {
   t_real const gamma = 0.0001;
 
   std::string tf_model_path =
-      purify::notinstalled::data_directory() + "/models/snr_15_model_dynamic.onnx";
+      purify::notinstalled::models_directory() + "/snr_15_model_dynamic.onnx";
 
   auto const fb = factory::fb_factory<sopt::algorithm::ImagingForwardBackward<t_complex>>(
       factory::algo_distribution::serial, measurements_transform, wavelets, uv_data, sigma, beta,
@@ -287,9 +287,9 @@ TEST_CASE("onnx_fb_factory") {
   t_real const gamma = 0.0001;
 
   std::string const prior_path =
-      purify::notinstalled::data_directory() + "/models/example_cost_dynamic_CRR_sigma_5_t_5.onnx";
+      purify::notinstalled::models_directory() + "/example_cost_dynamic_CRR_sigma_5_t_5.onnx";
   std::string const prior_gradient_path =
-      purify::notinstalled::data_directory() + "/models/example_grad_dynamic_CRR_sigma_5_t_5.onnx";
+      purify::notinstalled::models_directory() + "/example_grad_dynamic_CRR_sigma_5_t_5.onnx";
   std::shared_ptr<sopt::ONNXDifferentiableFunc<t_complex>> diff_function =
       std::make_shared<sopt::ONNXDifferentiableFunc<t_complex>>(
           prior_path, prior_gradient_path, sigma, 20, 5e4, *measurements_transform);

diff --git a/cpp/tests/directories.in.h b/cpp/tests/directories.in.h
@@ -8,6 +8,9 @@ namespace notinstalled {
 
 //! Holds data and such
 inline std::string data_directory() { return "@CMAKE_INSTALL_PREFIX@/data"; }
+//! Holds TF models
+inline std::string models_directory() { return "@CMAKE_INSTALL_PREFIX@/models"; }
+
 //! Holds data and such
 inline std::string data_filename(std::string const &filename) {
   return data_directory() + "/" + filename;

diff --git a/...example_cost_dynamic_CRR_sigma_5_t_5.onnx → ...example_cost_dynamic_CRR_sigma_5_t_5.onnx b/...example_cost_dynamic_CRR_sigma_5_t_5.onnx → ...example_cost_dynamic_CRR_sigma_5_t_5.onnx
diff --git a/...example_grad_dynamic_CRR_sigma_5_t_5.onnx → ...example_grad_dynamic_CRR_sigma_5_t_5.onnx b/...example_grad_dynamic_CRR_sigma_5_t_5.onnx → ...example_grad_dynamic_CRR_sigma_5_t_5.onnx
diff --git a/data/models/snr_15_model_dynamic.onnx → models/snr_15_model_dynamic.onnx b/data/models/snr_15_model_dynamic.onnx → models/snr_15_model_dynamic.onnx