[experimental] Dino for LR prediction

.gitignore

@@ -125,3 +125,6 @@ data
 results
 
 !examples/benchmarks/compression/results/
+
+wandb/
+temp/

examples/image_fitting.py

@@ -11,7 +11,18 @@
 from torch import Tensor, optim
 
 from gsplat import rasterization, rasterization_2dgs
+import random
 
+def seed_everything(seed: int):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    np.random.seed(seed)
+    random.seed(seed)
+
+seed_everything(42)
 
 class SimpleTrainer:
     """Trains random gaussians to fit an image."""
@@ -79,8 +90,10 @@ def train(
         iterations: int = 1000,
         lr: float = 0.01,
         save_imgs: bool = False,
+        save_path: str = None,
         model_type: Literal["3dgs", "2dgs"] = "3dgs",
     ):
+        losses = []
         optimizer = optim.Adam(
             [self.rgbs, self.means, self.scales, self.opacities, self.quats], lr
         )
@@ -100,9 +113,12 @@ def train(
             rasterize_fnc = rasterization
         elif model_type == "2dgs":
             rasterize_fnc = rasterization_2dgs
-
-        for iter in range(iterations):
+        verbose=False
+        from tqdm import tqdm
+        pbar = tqdm(range(iterations))
+        for iter in pbar:
             start = time.time()
+
 
             renders = rasterize_fnc(
                 self.means,
@@ -120,13 +136,17 @@ def train(
             torch.cuda.synchronize()
             times[0] += time.time() - start
             loss = mse_loss(out_img, self.gt_image)
+            if iter == 0:
+                init_loss = loss
             optimizer.zero_grad()
             start = time.time()
             loss.backward()
             torch.cuda.synchronize()
             times[1] += time.time() - start
             optimizer.step()
-            print(f"Iteration {iter + 1}/{iterations}, Loss: {loss.item()}")
+            losses.append(loss.item())
+            if verbose:
+                print(f"Iteration {iter + 1}/{iterations}, Loss: {loss.item()}")
 
             if save_imgs and iter % 5 == 0:
                 frames.append((out_img.detach().cpu().numpy() * 255).astype(np.uint8))
@@ -143,11 +163,21 @@ def train(
                 duration=5,
                 loop=0,
             )
-        print(f"Total(s):\nRasterization: {times[0]:.3f}, Backward: {times[1]:.3f}")
-        print(
+
+        if verbose:
+            print(f"Total(s):\nRasterization: {times[0]:.3f}, Backward: {times[1]:.3f}")
+            print(
             f"Per step(s):\nRasterization: {times[0]/iterations:.5f}, Backward: {times[1]/iterations:.5f}"
         )
 
+        final_img = (out_img.detach().cpu().numpy() * 255).astype(np.uint8)
+        if save_path:
+
+            import matplotlib.pyplot as plt
+            plt.imsave(save_path, final_img)
+
+        return losses, final_img
+
 
 def image_path_to_tensor(image_path: Path):
     import torchvision.transforms as transforms

examples/tile_sweep.sh

@@ -0,0 +1,87 @@
+# #!/bin/bash
+
+# # Function to run the tile_trainer.py script
+# run_tile_trainer() {
+#     iterations=$1
+#     num_points=$2
+#     gpu=$3
+#     echo "Running on GPU $gpu: iterations=$iterations, num_points=$num_points"
+#     CUDA_VISIBLE_DEVICES=$gpu python examples/tile_trainer.py --weights 0.25 0.25 0.25 0.25 --iterations $iterations --num_points $num_points
+# }
+
+# # Function to find the next available GPU
+# find_available_gpu() {
+#     for gpu in 0 1 2; do
+#         if [ $(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i $gpu) -lt 80 ]; then
+#             echo $gpu
+#             return
+#         fi
+#     done
+#     echo "-1"
+# }
+
+# # Generate all combinations of iterations and num_points
+# combinations=()
+# for iterations in $(seq 500 500 9500); do
+#     for num_points in $(seq 100 100 4900); do
+#         combinations+=("$iterations $num_points")
+#     done
+# done
+
+# # Process all combinations
+# for combination in "${combinations[@]}"; do
+#     read iterations num_points <<< "$combination"
+
+#     # Find an available GPU
+#     while true; do
+#         gpu=$(find_available_gpu)
+#         if [ "$gpu" != "-1" ]; then
+#             run_tile_trainer $iterations $num_points $gpu &
+#             sleep 1  # Short delay to allow the job to start
+#             break
+#         fi
+#         sleep 5  # Wait before checking GPUs again
+#     done
+# done
+
+# # Wait for all background processes to finish
+# wait
+
+# echo "All jobs completed."
+
+
+
+
+#!/bin/bash
+
+# Initialize an array to track GPU processes
+gpu_pids=()
+
+# Set the ranges for iterations and num_points
+for iterations in $(seq 500 500 1500); do
+  for num_points in $(seq 100 100 300); do
+    # Calculate the GPU to assign (cycles through 0, 1, 2 for 3 GPUs)
+    gpu=$(( (counter % 3) ))
+
+    # Wait for the GPU to finish its previous task if it's still running
+    if [ ! -z "${gpu_pids[$gpu]}" ]; then
+      wait ${gpu_pids[$gpu]}
+    fi
+
+    # Run the tile_trainer.py command and store the process ID (PID)
+    CUDA_VISIBLE_DEVICES=$gpu python examples/tile_trainer.py --random --weights 0.4 0.1 0.1 0.4 \
+      --iterations $iterations --num_points $num_points &
+
+    # Store the PID of the background process
+    gpu_pids[$gpu]=$!
+
+    # Increment the counter
+    counter=$((counter + 1))
+
+    # Small delay to avoid overwhelming the system
+    sleep 0.2
+  done
+done
+
+# Wait for all remaining GPU jobs to finish
+wait

merge_dicts.ipynb

@@ -0,0 +1,214 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import json\n",
+    "\n",
+    "dict1_path = \"/secondary/home/aayushg/rl_gsplat/gsplat/src/ppo/lr_predictor/small_mipnerf_lr_losses_0.025_startlr_1000_iterations100000_points_10_trials_horns_DJI_20200223_163035_787.png.json\"\n",
+    "dict2_path = \"/secondary/home/aayushg/rl_gsplat/gsplat/src/ppo/lr_predictor/small_mipnerf_lr_losses_0.035_startlr_1000_iterations100000_points_10_trials_horns_DJI_20200223_163035_787.png.json\"\n",
+    "dict3_path = \"/secondary/home/aayushg/rl_gsplat/gsplat/src/ppo/lr_predictor/small_mipnerf_lr_losses_backup.json\"\n",
+    "\n",
+    "with open(dict1_path, \"r\") as f:\n",
+    "    dict1 = json.load(f)\n",
+    "\n",
+    "with open(dict2_path, \"r\") as f:\n",
+    "    dict2 = json.load(f)\n",
+    "\n",
+    "with open(dict3_path, \"r\") as f:\n",
+    "    dict3 = json.load(f)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "dict_keys(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12', '13', '14', '15', '16', '17', '18', '19', '20', '21', '22', '23', '24', '25', '26', '27', '28', '29', '30', '31'])"
+      ]
+     },
+     "execution_count": 18,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "dict1.keys()\n",
+    "dict2.keys()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for img_idx in dict1.keys():\n",
+    "    dict1_lr_losses = dict1[img_idx][\"lr_losses\"]\n",
+    "    dict2_lr_losses = dict2[img_idx][\"lr_losses\"]4\n",
+    "    for lr in dict1_lr_losses.keys():\n",
+    "        dict3[img_idx][\"lr_losses\"][lr] = dict1_lr_losses[lr]\n",
+    "    for lr in dict2_lr_losses.keys():\n",
+    "        dict3[img_idx][\"lr_losses\"][lr] = dict2_lr_losses[lr]\n",
+    "\n",
+    "with open(dict3_path + \"_new_merged\", \"w\") as f:\n",
+    "    json.dump(dict3, f)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "dict_keys(['0.005', '0.006', '0.007', '0.008', '0.009', '0.01', '0.011', '0.012', '0.013', '0.014', '0.015', '0.016', '0.017', '0.018', '0.019', '0.02', '0.021', '0.022', '0.023', '0.024', '0.025', '0.026', '0.027', '0.028', '0.029', '0.03', '0.031', '0.032', '0.033', '0.034', '0.035', '0.036', '0.037', '0.038', '0.039', '0.04', '0.041', '0.042', '0.043', '0.044'])"
+      ]
+     },
+     "execution_count": 22,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "dict3['31']['lr_losses'].keys()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "dict_keys(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12', '13', '14', '15', '16', '17', '18'])"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "dict1.keys()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "dict_keys(['lr_losses', 'img_name'])"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "dict1[\"0\"].keys()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "dict_keys(['0.025', '0.026', '0.027', '0.028', '0.029', '0.03', '0.031', '0.032', '0.033', '0.034'])"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "dict1[\"0\"][\"lr_losses\"].keys()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "dict_keys(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12', '13', '14', '15', '16', '17', '18'])"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "dict2.keys()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "dict_keys(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12', '13', '14', '15', '16', '17', '18', '19', '20', '21', '22', '23', '24', '25', '26', '27', '28', '29', '30', '31'])"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "dict3.keys()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "gsplat",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}