Snapshot extension is now ".ingp" and is zip compressed

Author: Tom94

.gitignore

@@ -12,6 +12,7 @@
 /*.dll
 /*.exe
 /*.json
+*.ingp
 *.msgpack
 *.training
 __pycache__

.gitmodules

@@ -22,3 +22,9 @@
 [submodule "dependencies/dlss"]
 	path = dependencies/dlss
 	url = https://github.com/NVIDIA/DLSS
+[submodule "dependencies/zstr"]
+	path = dependencies/zstr
+	url = https://github.com/Tom94/zstr
+[submodule "dependencies/zlib"]
+	path = dependencies/zlib
+	url = https://github.com/Tom94/zlib

CMakeLists.txt

@@ -193,6 +193,26 @@ if (Python_FOUND)
 	add_subdirectory("dependencies/pybind11")
 endif()
 
+# Compile zlib (only on Windows)
+if (WIN32)
+	set(ZLIB_USE_STATIC_LIBS ON CACHE BOOL " " FORCE)
+	set(ZLIB_BUILD_STATIC_LIBS ON CACHE BOOL " " FORCE)
+	set(ZLIB_BUILD_SHARED_LIBS OFF CACHE BOOL " " FORCE)
+	set(SKIP_INSTALL_ALL ON CACHE BOOL " " FORCE)
+	add_subdirectory("dependencies/zlib")
+	set_property(TARGET zlibstatic PROPERTY FOLDER "dependencies")
+
+	set(ZLIB_INCLUDE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/dependencies/zlib" CACHE PATH " " FORCE)
+	set(ZLIB_LIBRARY zlibstatic)
+
+	include_directories(${ZLIB_INCLUDE_DIR} "${CMAKE_CURRENT_BINARY_DIR}/dependencies/zlib")
+
+	list(APPEND NGP_LIBRARIES zlibstatic)
+endif()
+
+add_subdirectory("dependencies/zstr")
+list(APPEND NGP_LIBRARIES zstr::zstr)
+
 
 ###############################################################################
 # Program

dependencies/zlib

@@ -412,7 +412,7 @@ class Testbed {
 	void set_fov(float val) ;
 	Eigen::Vector2f fov_xy() const ;
 	void set_fov_xy(const Eigen::Vector2f& val);
-	void save_snapshot(const std::string& filepath_string, bool include_optimizer_state);
+	void save_snapshot(const std::string& filepath_string, bool include_optimizer_state, bool compress);
 	void load_snapshot(const std::string& filepath_string);
 	CameraKeyframe copy_camera_to_keyframe() const;
 	void set_camera_from_keyframe(const CameraKeyframe& k);
@@ -464,6 +464,7 @@ class Testbed {
 	bool m_gather_histograms = false;
 
 	bool m_include_optimizer_state_in_snapshot = false;
+	bool m_compress_snapshot = true;
 	bool m_render_ground_truth = false;
 	EGroundTruthRenderMode m_ground_truth_render_mode = EGroundTruthRenderMode::Shade;
 	float m_ground_truth_alpha = 1.0f;
@@ -831,7 +832,7 @@ class Testbed {
 		char cam_path_path[MAX_PATH_LEN] = "cam.json";
 		char extrinsics_path[MAX_PATH_LEN] = "extrinsics.json";
 		char mesh_path[MAX_PATH_LEN] = "base.obj";
-		char snapshot_path[MAX_PATH_LEN] = "base.msgpack";
+		char snapshot_path[MAX_PATH_LEN] = "base.ingp";
 		char video_path[MAX_PATH_LEN] = "video.mp4";
 	} m_imgui;
 

dependencies/zstr

@@ -5846,7 +5846,7 @@
       ],
       "source": [
         "train_steps = 2000  #@param {type:\"integer\"}\n",
-        "snapshot_path = os.path.join(scene_path, f\"{train_steps}.msgpack\")\n",
+        "snapshot_path = os.path.join(scene_path, f\"{train_steps}.ingp\")\n",
         "!python ./scripts/run.py {scene_path} --n_steps {train_steps} --save_snapshot {snapshot_path}"
       ]
     },
@@ -5863,7 +5863,7 @@
         "\n",
         "Example command:\n",
         "```\n",
-        "./build/instant-ngp /data/nerf/fox/2000.msgpack\n",
+        "./build/instant-ngp /data/nerf/fox/2000.ingp\n",
         "```\n",
         "\n",
         "After you're done, **upload `base_cam.json` to the root folder of your scene.**"

include/neural-graphics-primitives/testbed.h

@@ -33,10 +33,10 @@ def parse_args():
 	parser.add_argument("--mode", default="", type=str, help=argparse.SUPPRESS) # deprecated
 	parser.add_argument("--network", default="", help="Path to the network config. Uses the scene's default if unspecified.")
 
-	parser.add_argument("--load_snapshot", "--snapshot", default="", help="Load this snapshot before training. recommended extension: .msgpack")
-	parser.add_argument("--save_snapshot", default="", help="Save this snapshot after training. recommended extension: .msgpack")
+	parser.add_argument("--load_snapshot", "--snapshot", default="", help="Load this snapshot before training. recommended extension: .ingp/.msgpack")
+	parser.add_argument("--save_snapshot", default="", help="Save this snapshot after training. recommended extension: .ingp/.msgpack")
 
-	parser.add_argument("--nerf_compatibility", action="store_true", help="Matches parameters with original NeRF. Can cause slowness and worse results on some scenes.")
+	parser.add_argument("--nerf_compatibility", action="store_true", help="Matches parameters with original NeRF. Can cause slowness and worse results on some scenes, but helps with high PSNR on synthetic scenes.")
 	parser.add_argument("--test_transforms", default="", help="Path to a nerf style transforms json from which we will compute PSNR.")
 	parser.add_argument("--near_distance", default=-1, type=float, help="Set the distance from the camera at which training rays start for nerf. <0 means use ngp default")
 	parser.add_argument("--exposure", default=0.0, type=float, help="Controls the brightness of the image. Positive numbers increase brightness, negative numbers decrease it.")

notebooks/instant_ngp.ipynb

@@ -222,7 +222,7 @@ def setup_colored_sdf(testbed, scene, softshadow=True):
 	testbed.scale = testbed.scale * 1.13
 
 def default_snapshot_filename(scene_info):
-	filename = "base.msgpack"
+	filename = f"base.ingp"
 	if scene_info["dataset"]:
 		filename = f"{os.path.splitext(scene_info['dataset'])[0]}_{filename}"
 	return os.path.join(scene_info["data_dir"], filename)

scripts/run.py

@@ -41,6 +41,8 @@
 #include <stb_image/stb_image.h>
 #include <stb_image/stb_image_write.h>
 
+#include <zstr.hpp>
+
 #include <fstream>
 #include <set>
 #include <unordered_set>
@@ -86,7 +88,7 @@ NGP_NAMESPACE_BEGIN
 
 std::atomic<size_t> g_total_n_bytes_allocated{0};
 
-json merge_parent_network_config(const json &child, const fs::path &child_filename) {
+json merge_parent_network_config(const json& child, const fs::path& child_filename) {
 	if (!child.contains("parent")) {
 		return child;
 	}
@@ -121,7 +123,7 @@ void Testbed::update_imgui_paths() {
 	snprintf(m_imgui.cam_path_path, sizeof(m_imgui.cam_path_path), "%s", get_filename_in_data_path_with_suffix(m_data_path, m_network_config_path, "_cam.json").c_str());
 	snprintf(m_imgui.extrinsics_path, sizeof(m_imgui.extrinsics_path), "%s", get_filename_in_data_path_with_suffix(m_data_path, m_network_config_path, "_extrinsics.json").c_str());
 	snprintf(m_imgui.mesh_path, sizeof(m_imgui.mesh_path), "%s", get_filename_in_data_path_with_suffix(m_data_path, m_network_config_path, ".obj").c_str());
-	snprintf(m_imgui.snapshot_path, sizeof(m_imgui.snapshot_path), "%s", get_filename_in_data_path_with_suffix(m_data_path, m_network_config_path, ".msgpack").c_str());
+	snprintf(m_imgui.snapshot_path, sizeof(m_imgui.snapshot_path), "%s", get_filename_in_data_path_with_suffix(m_data_path, m_network_config_path, ".ingp").c_str());
 	snprintf(m_imgui.video_path, sizeof(m_imgui.video_path), "%s", get_filename_in_data_path_with_suffix(m_data_path, m_network_config_path, "_video.mp4").c_str());
 }
 
@@ -215,22 +217,28 @@ fs::path Testbed::find_network_config(const fs::path& network_config_path) {
 }
 
 json Testbed::load_network_config(const fs::path& network_config_path) {
-	bool is_snapshot = equals_case_insensitive(network_config_path.extension(), "msgpack");
+	bool is_snapshot = equals_case_insensitive(network_config_path.extension(), "msgpack") || equals_case_insensitive(network_config_path.extension(), "ingp");
 	if (network_config_path.empty() || !network_config_path.exists()) {
 		throw std::runtime_error{fmt::format("Network {} '{}' does not exist.", is_snapshot ? "snapshot" : "config", network_config_path.str())};
 	}
 
 	tlog::info() << "Loading network " << (is_snapshot ? "snapshot" : "config") << " from: " << network_config_path;
 
 	json result;
-	if (equals_case_insensitive(network_config_path.extension(), "json")) {
+	if (is_snapshot) {
+		if (equals_case_insensitive(network_config_path.extension(), "ingp")) {
+			// zstr::ifstream applies zlib compression.
+			zstr::ifstream f{network_config_path.str(), std::ios::in | std::ios::binary};
+			result = json::from_msgpack(f);
+		} else {
+			std::ifstream f{network_config_path.str(), std::ios::in | std::ios::binary};
+			result = json::from_msgpack(f);
+		}
+		// we assume parent pointers are already resolved in snapshots.
+	} else if (equals_case_insensitive(network_config_path.extension(), "json")) {
 		std::ifstream f{network_config_path.str()};
 		result = json::parse(f, nullptr, true, true);
 		result = merge_parent_network_config(result, network_config_path);
-	} else if (equals_case_insensitive(network_config_path.extension(), "msgpack")) {
-		std::ifstream f{network_config_path.str(), std::ios::in | std::ios::binary};
-		result = json::from_msgpack(f);
-		// we assume parent pointers are already resolved in snapshots.
 	}
 
 	return result;
@@ -289,7 +297,7 @@ void Testbed::load_file(const std::string& file_path) {
 		return;
 	}
 
-	if (ends_with_case_insensitive(file_path, ".msgpack")) {
+	if (ends_with_case_insensitive(file_path, ".ingp") || ends_with_case_insensitive(file_path, ".msgpack")) {
 		load_snapshot(file_path);
 		return;
 	}
@@ -1218,7 +1226,7 @@ void Testbed::imgui() {
 		ImGui::Text("Snapshot");
 		ImGui::SameLine();
 		if (ImGui::Button("Save")) {
-			save_snapshot(m_imgui.snapshot_path, m_include_optimizer_state_in_snapshot);
+			save_snapshot(m_imgui.snapshot_path, m_include_optimizer_state_in_snapshot, m_compress_snapshot);
 		}
 		ImGui::SameLine();
 		if (ImGui::Button("Load")) {
@@ -1237,6 +1245,16 @@ void Testbed::imgui() {
 		ImGui::SameLine();
 		ImGui::Checkbox("w/ optimizer state", &m_include_optimizer_state_in_snapshot);
 		ImGui::InputText("File##Snapshot file path", m_imgui.snapshot_path, sizeof(m_imgui.snapshot_path));
+		ImGui::SameLine();
+
+		bool can_compress = ends_with_case_insensitive(m_imgui.snapshot_path, ".ingp");
+
+		if (!can_compress) {
+			ImGui::BeginDisabled();
+			m_compress_snapshot = false;
+		}
+		ImGui::Checkbox("Compress", &m_compress_snapshot);
+		if (!can_compress) ImGui::EndDisabled();
 	}
 
 	if (m_testbed_mode == ETestbedMode::Nerf || m_testbed_mode == ETestbedMode::Sdf) {
@@ -3481,7 +3499,7 @@ void Testbed::gather_histograms() {
 // Increment this number when making a change to the snapshot format
 static const size_t SNAPSHOT_FORMAT_VERSION = 1;
 
-void Testbed::save_snapshot(const std::string& filepath_string, bool include_optimizer_state) {
+void Testbed::save_snapshot(const std::string& filepath_string, bool include_optimizer_state, bool compress) {
 	fs::path filepath = filepath_string;
 	m_network_config["snapshot"] = m_trainer->serialize(include_optimizer_state);
 
@@ -3516,8 +3534,16 @@ void Testbed::save_snapshot(const std::string& filepath_string, bool include_opt
 	}
 
 	m_network_config_path = filepath;
-	std::ofstream f(m_network_config_path.str(), std::ios::out | std::ios::binary);
-	json::to_msgpack(m_network_config, f);
+	if (equals_case_insensitive(m_network_config_path.extension(), "ingp")) {
+		// zstr::ofstream applies zlib compression.
+		zstr::ofstream f{m_network_config_path.str(), std::ios::out | std::ios::binary, compress ? Z_DEFAULT_COMPRESSION : Z_NO_COMPRESSION};
+		json::to_msgpack(m_network_config, f);
+	} else {
+		std::ofstream f{m_network_config_path.str(), std::ios::out | std::ios::binary};
+		json::to_msgpack(m_network_config, f);
+	}
+
+	tlog::success() << "Saved snapshot '" << filepath_string << "'";
 }
 
 void Testbed::load_snapshot(const std::string& filepath_string) {
@@ -3526,10 +3552,7 @@ void Testbed::load_snapshot(const std::string& filepath_string) {
 		throw std::runtime_error{fmt::format("File '{}' does not contain a snapshot.", filepath_string)};
 	}
 
-	m_network_config_path = filepath_string;
-
 	const auto& snapshot = config["snapshot"];
-
 	if (snapshot.value("version", 0) < SNAPSHOT_FORMAT_VERSION) {
 		throw std::runtime_error{"Snapshot uses an old format and can not be loaded."};
 	}