Implement fused SSIM (#30)

* Implement fused ssim

* Remove Flux

* Update tests & cleanup deps

Author: pxl-th

LocalPreferences.toml

@@ -6,13 +6,11 @@ authors = ["Anton Smirnov <[email protected]>"]
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
-BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CImGui = "5d785b6c-b76f-510e-a07c-3070796c7e87"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
-Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 GLFW = "f7f18e0c-5ee9-5ccd-a5bf-e8befd85ed98"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
@@ -29,7 +27,6 @@ NearestNeighbors = "b8a86587-4115-5ab1-83bc-aa920d37bbce"
 NerfUtils = "99c1d5ce-7c61-4a25-a107-a5ade2e2a8e4"
 NeuralGraphicsGL = "263f7e6d-e369-49af-a86e-c85638573b76"
 PlyIO = "42171d58-473b-503a-8d5f-782019eb09ec"
-Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Quaternions = "94ee1d12-ae83-5a48-8b1c-48b8ff168ae0"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Rotations = "6038ab10-8711-5258-84ad-4b1120ba62dc"
@@ -42,23 +39,20 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
 [extensions]
 GaussianSplattingAMDGPUExt = "AMDGPU"
-GaussianSplattingCUDAExt = ["CUDA", "cuDNN"]
+GaussianSplattingCUDAExt = "CUDA"
 
 [compat]
 AMDGPU = "2"
 Adapt = "4"
 BSON = "0.3"
-BenchmarkTools = "1"
 CImGui = "6"
 CUDA = "5"
 ChainRulesCore = "1"
 Distributions = "0.25"
 FileIO = "1.16"
-Flux = "0.16.2"
 GLFW = "3.4"
 GPUArrays = "11.2.1"
 GPUArraysCore = "0.2"
@@ -74,12 +68,10 @@ NearestNeighbors = "0.4"
 NerfUtils = "0.2"
 NeuralGraphicsGL = "0.5"
 PlyIO = "1.2"
-Preferences = "1.4"
 Quaternions = "0.7"
 Rotations = "1.7"
 SIMD = "3.6"
 StaticArrays = "1.9"
 VideoIO = "1.1"
 Zygote = "=0.7.3"
-cuDNN = "1.3"
 julia = "1.10"

Project.toml

@@ -25,17 +25,14 @@ GaussianSplatting.jl comes with a GUI application to train & view the gaussians.
 1. Add necessary packages:
    ```julia
    ] add AMDGPU      # for AMD GPU
-   ] add CUDA, cuDNN # for Nvidia GPU
-   ] add Flux
+   ] add CUDA # for Nvidia GPU
    ```
 
 2. Run:
    ```julia
-   julia> using AMDGPU      # for AMD GPU
-   julia> using CUDA, cuDNN # for Nvidia GPU
-   julia> using Flux, GaussianSplatting
-
-   julia> GaussianSplatting.gui("path-to-colmap-dataset-directory"; scale=1)
+   julia> using AMDGPU; kab = ROCBackend() # for AMD GPU
+   julia> using CUDA; kab = CUDABackend()  # for Nvidia GPU
+   julia> GaussianSplatting.gui(kab, "path-to-colmap-dataset-directory"; scale=1)
    ```
 
 ## Viewer mode
@@ -44,27 +41,12 @@ Once you've trained a model and saved it to `.bson` file you can open it
 in a viewer-only mode by providing its path.
 
 ```julia
-julia> GaussianSplatting.gui("path-to-checkpoint.bson")
+julia> GaussianSplatting.gui(kab, "path-to-checkpoint.bson")
 ```
 
 Alternative, you can load a checkpoint
 in a training mode (see **Usage** section) using "Save/Load" tab.
 
-## GPU selection
-
-This is required only the first time per the environment.
-After selecting GPU backend, restart Julia REPL.
-
-- AMD GPU:
-  ```julia
-  julia> Flux.gpu_backend!("AMDGPU")
-  ```
-
-- Nvidia GPU:
-  ```julia
-  julia> Flux.gpu_backend!("CUDA")
-  ```
-
 ## References
 
 - 3D Gaussian Splatting for Real-Time Radiance Field Rendering: <https://arxiv.org/abs/2308.04079>

docs/src/index.md

@@ -2,7 +2,6 @@ module GaussianSplattingCUDAExt
 
 # using Adapt
 using CUDA
-# using cuDNN
 # using KernelAbstractions
 using GaussianSplatting
 # using PrecompileTools

ext/GaussianSplattingCUDAExt/GaussianSplattingCUDAExt.jl

@@ -17,7 +17,6 @@ using Random
 using Rotations
 using StaticArrays
 using Statistics
-using Preferences
 using ImageCore
 using ImageIO
 using ImageTransformations
@@ -33,8 +32,8 @@ using GLFW
 import CImGui.lib as iglib
 
 import BSON
+import ChainRulesCore as CRC
 import NNlib
-import Flux
 import ImageFiltering
 import KernelAbstractions as KA
 import NerfUtils as NU
@@ -56,7 +55,7 @@ _as_T(T, x) = reinterpret(T, reshape(x, :))
 
 include("simd.jl")
 include("utils.jl")
-include("metrics.jl")
+include("fused_ssim.jl")
 include("camera.jl")
 include("camera_opt.jl")
 include("dataset.jl")
@@ -67,9 +66,6 @@ include("rasterization/rasterizer.jl")
 include("training.jl")
 include("gui/gui.jl")
 
-# Hacky way to get KA.Backend.
-gpu_backend() = get_backend(Flux.gpu(Array{Int}(undef, 0)))
-
 base_array_type(backend) = error("Not implemented for backend: `$backend`.")
 
 allocate_pinned(kab, T, shape) = error("Pinned memory not supported for `$kab`.")
@@ -78,8 +74,7 @@ unpin_memory(x) = error("Unpinning memory is not supported for `$(typeof(x))`.")
 
 use_ak(kab) = false
 
-function main(dataset_path::String; scale::Int, save_path::Maybe{String} = nothing)
-    kab = gpu_backend()
+function main(kab, dataset_path::String; scale::Int, save_path::Maybe{String} = nothing)
     @info "Using `$kab` GPU backend."
 
     dataset = ColmapDataset(kab, dataset_path;
@@ -151,7 +146,7 @@ function main(dataset_path::String; scale::Int, save_path::Maybe{String} = nothi
     return
 end
 
-function gui(path::String; scale::Maybe{Int} = nothing, fullscreen::Bool = false)
+function gui(kab, path::String; scale::Maybe{Int} = nothing, fullscreen::Bool = false)
     ispath(path) || error("Path does not exist: `$path`.")
 
     viewer_mode = endswith(path, ".bson") || endswith(path, ".ply")
@@ -166,7 +161,6 @@ function gui(path::String; scale::Maybe{Int} = nothing, fullscreen::Bool = false
         (1024, 1024, true)
 
     gui = if viewer_mode
-        kab = gpu_backend()
         if endswith(path, ".bson")
             θ = BSON.load(path)
             gaussians = GaussianModel(kab)
@@ -179,9 +173,9 @@ function gui(path::String; scale::Maybe{Int} = nothing, fullscreen::Bool = false
             camera = Camera(; fx=fov, fy=fov, width, height=width)
         end
 
-        GSGUI(gaussians, camera; width, height, fullscreen, resizable)
+        GSGUI(kab, gaussians, camera; width, height, fullscreen, resizable)
     else
-        GSGUI(path, scale; width, height, fullscreen, resizable)
+        GSGUI(kab, path, scale; width, height, fullscreen, resizable)
     end
     gui |> launch!
     return