Skip to content

Commit 33b4a66

Browse files
committed
Merge branch 'dev_compression_simulation' into main
2 parents ee93e67 + 179d3a3 commit 33b4a66

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

41 files changed

+8024
-378
lines changed

.github/workflows/publish.yml

Lines changed: 0 additions & 111 deletions
This file was deleted.

.gitignore

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,8 @@ target/
7777

7878
# Jupyter Notebook
7979
.ipynb_checkpoints
80+
.ipynb
81+
*.ipynb
8082

8183
# pyenv
8284
.python-version
@@ -124,7 +126,17 @@ compile_commands.json
124126
data
125127
results
126128

129+
third_party
130+
# gscodec
131+
setup_gscodec.py
132+
Readme_GSCodec.md
133+
134+
figs
135+
stats
136+
temp
137+
127138
!examples/benchmarks/compression/results/
139+
.github/
128140

129141
setup_gscodec.py
130142
*.ipynb
Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
# --------------- training settings -------------------- #
2+
3+
SCENE_DIR="data/db"
4+
5+
SCENE_LIST="drjohnson playroom" #
6+
7+
# # 0.36M GSs
8+
# RESULT_DIR="results/benchmark_db_mcmc_0_36M_png_compression"
9+
# CAP_MAX=360000
10+
11+
# # 0.49M GSs
12+
RESULT_DIR="results/benchmark_db_mcmc_0_49M_png_compression"
13+
CAP_MAX=490000
14+
15+
# 1M GSs
16+
# RESULT_DIR="results/benchmark_db_mcmc_1M_png_compression"
17+
# CAP_MAX=1000000
18+
19+
# # 4M GSs
20+
# RESULT_DIR="results/benchmark_db_mcmc_4M_png_compression"
21+
# CAP_MAX=4000000
22+
23+
# Override default values if provided as arguments
24+
# [ ! -z "$1" ] && RESULT_DIR="$1"
25+
# [ ! -z "$2" ] && CAP_MAX="$2"
26+
27+
RD_LAMBDA=0.01
28+
29+
# ----------------- Training Setting-------------- #
30+
31+
# ----------------- Args ------------------------- #
32+
33+
if [ ! -z "$1" ]; then
34+
RD_LAMBDA="$1"
35+
RESULT_DIR="results/Ours_DB_rd_lambda_${RD_LAMBDA}"
36+
fi
37+
38+
# ----------------- Args ------------------------- #
39+
40+
# ----------------- Main Job --------------------- #
41+
run_single_scene() {
42+
local GPU_ID=$1
43+
local SCENE=$2
44+
45+
echo "Running $SCENE on GPU: $GPU_ID"
46+
47+
# train without eval
48+
CUDA_VISIBLE_DEVICES=$GPU_ID python simple_trainer.py mcmc --eval_steps -1 --disable_viewer --data_factor 1 \
49+
--strategy.cap-max $CAP_MAX \
50+
--data_dir $SCENE_DIR/$SCENE/ \
51+
--result_dir $RESULT_DIR/$SCENE/ \
52+
--compression_sim \
53+
--entropy_model_opt \
54+
--rd_lambda $RD_LAMBDA \
55+
--shN_ada_mask_opt
56+
# --opacity_reg 0.001
57+
58+
59+
eval: use vgg for lpips to align with other benchmarks
60+
CUDA_VISIBLE_DEVICES=$GPU_ID python simple_trainer.py mcmc --disable_viewer --data_factor 1 \
61+
--strategy.cap-max $CAP_MAX \
62+
--data_dir $SCENE_DIR/$SCENE/ \
63+
--result_dir $RESULT_DIR/$SCENE/ \
64+
--lpips_net vgg \
65+
--compression png \
66+
--ckpt $RESULT_DIR/$SCENE/ckpts/ckpt_29999_rank0.pt
67+
68+
}
69+
# ----------------- Main Job --------------------- #
70+
71+
GPU_LIST=(6 7)
72+
GPU_COUNT=${#GPU_LIST[@]}
73+
74+
SCENE_IDX=-1
75+
76+
for SCENE in $SCENE_LIST;
77+
do
78+
SCENE_IDX=$((SCENE_IDX + 1))
79+
{
80+
run_single_scene ${GPU_LIST[$SCENE_IDX]} $SCENE
81+
} &
82+
83+
done
84+
85+
# ----------------- Experiment Loop -------------- #
86+
87+
# Wait for finishing the jobs across all scenes
88+
wait
89+
echo "All scenes finished."
90+
91+
# Zip the compressed files and summarize the stats
92+
if command -v zip &> /dev/null
93+
then
94+
echo "Zipping results"
95+
python benchmarks/compression/summarize_stats.py --results_dir $RESULT_DIR --scenes $SCENE_LIST
96+
else
97+
echo "zip command not found, skipping zipping"
98+
fi
Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
# ------------------------- Settings ------------------------- #
2+
SCENE_DIR="data/360_v2"
3+
# eval all 9 scenes for benchmarking
4+
SCENE_LIST="bicycle bonsai counter flowers garden kitchen room stump treehill"
5+
#
6+
7+
# # 0.36M GSs
8+
# RESULT_DIR="results/benchmark_mcmc_0_36M_png_compression"
9+
# CAP_MAX=360000
10+
11+
# # 0.49M GSs
12+
# RESULT_DIR="results/benchmark_mcmc_0_49M_png_compression"
13+
# CAP_MAX=490000
14+
15+
# 1M GSs
16+
RESULT_DIR="results/Ours_MipNeRF360"
17+
CAP_MAX=1000000
18+
19+
# # 4M GSs
20+
# RESULT_DIR="results/benchmark_mcmc_4M_png_compression"
21+
# CAP_MAX=4000000
22+
23+
RD_LAMBDA=0.01
24+
25+
GPU_LIST="0 1 2 3 4"
26+
# ------------------------- Settings ------------------------- #
27+
28+
# ------------------------- Args ----------------------------- #
29+
30+
31+
if [ ! -z "$1" ]; then
32+
RD_LAMBDA="$1"
33+
RESULT_DIR="results/Ours_MipNeRF360_rd_lambda_${RD_LAMBDA}"
34+
fi
35+
36+
37+
# ------------------------- Args ----------------------------- #
38+
39+
# ------------------------- Experiment ----------------------- #
40+
41+
run_scene() {
42+
SCENE=$1
43+
GPU_ID=$2
44+
45+
if [ "$SCENE" = "bonsai" ] || [ "$SCENE" = "counter" ] || [ "$SCENE" = "kitchen" ] || [ "$SCENE" = "room" ]; then
46+
DATA_FACTOR=2
47+
else
48+
DATA_FACTOR=4
49+
fi
50+
51+
echo "Running $SCENE on GPU $GPU_ID"
52+
53+
# train without eval
54+
CUDA_VISIBLE_DEVICES=$GPU_ID python simple_trainer.py mcmc --eval_steps -1 --disable_viewer --data_factor $DATA_FACTOR \
55+
--strategy.cap-max $CAP_MAX \
56+
--data_dir data/360_v2/$SCENE/ \
57+
--result_dir $RESULT_DIR/$SCENE/ \
58+
--compression_sim \
59+
--entropy_model_opt \
60+
--rd_lambda $RD_LAMBDA \
61+
--shN_ada_mask_opt
62+
63+
# eval
64+
CUDA_VISIBLE_DEVICES=$GPU_ID python simple_trainer.py mcmc --disable_viewer --data_factor $DATA_FACTOR \
65+
--strategy.cap-max $CAP_MAX \
66+
--data_dir data/360_v2/$SCENE/ \
67+
--result_dir $RESULT_DIR/$SCENE/ \
68+
--lpips_net vgg \
69+
--compression png \
70+
--ckpt $RESULT_DIR/$SCENE/ckpts/ckpt_29999_rank0.pt
71+
}
72+
73+
# ------------------------- Experiment ----------------------- #
74+
75+
# ------------------------- Main ----------------------------- #
76+
77+
# 将场景列表和GPU列表转换为数组
78+
IFS=' ' read -r -a SCENE_ARRAY <<< "$SCENE_LIST"
79+
IFS=' ' read -r -a GPU_ARRAY <<< "$GPU_LIST"
80+
NUM_GPUS=${#GPU_ARRAY[@]}
81+
82+
# 追踪已经处理的场景索引
83+
current_scene_index=0
84+
total_scenes=${#SCENE_ARRAY[@]}
85+
86+
# 追踪每个GPU上运行的进程和场景
87+
declare -A gpu_pids
88+
declare -A gpu_scenes
89+
90+
# 初始分配场景给GPU
91+
for ((i=0; i<NUM_GPUS && i<total_scenes; i++)); do
92+
GPU_ID=${GPU_ARRAY[$i]}
93+
SCENE=${SCENE_ARRAY[$i]}
94+
run_scene $SCENE $GPU_ID &
95+
gpu_pids[$GPU_ID]=$!
96+
gpu_scenes[$GPU_ID]=$SCENE
97+
current_scene_index=$((i + 1))
98+
echo "Started $SCENE on GPU $GPU_ID"
99+
done
100+
101+
# 持续检查并分配新的场景
102+
while [ $current_scene_index -lt $total_scenes ]; do
103+
for GPU_ID in "${GPU_ARRAY[@]}"; do
104+
# 检查这个GPU上的进程是否完成
105+
if [ -n "${gpu_pids[$GPU_ID]}" ] && ! kill -0 ${gpu_pids[$GPU_ID]} 2>/dev/null; then
106+
echo "Finished ${gpu_scenes[$GPU_ID]} on GPU $GPU_ID"
107+
108+
# 分配新的场景给这个GPU
109+
SCENE=${SCENE_ARRAY[$current_scene_index]}
110+
run_scene $SCENE $GPU_ID &
111+
gpu_pids[$GPU_ID]=$!
112+
gpu_scenes[$GPU_ID]=$SCENE
113+
echo "Started $SCENE on GPU $GPU_ID"
114+
115+
current_scene_index=$((current_scene_index + 1))
116+
117+
# 如果所有场景都已分配,退出循环
118+
if [ $current_scene_index -eq $total_scenes ]; then
119+
break 2
120+
fi
121+
fi
122+
done
123+
sleep 10 # 检查间隔
124+
done
125+
126+
# 等待所有进程完成
127+
wait
128+
129+
echo "All scenes completed!"
130+
# ------------------------- Main ----------------------------- #
131+
132+
# ------------------------- Compression ---------------------- #
133+
134+
Zip the compressed files and summarize the stats
135+
if command -v zip &> /dev/null
136+
then
137+
echo "Zipping results"
138+
python benchmarks/compression/summarize_stats.py --results_dir $RESULT_DIR --scenes $SCENE_LIST
139+
else
140+
echo "zip command not found, skipping zipping"
141+
fi
142+
# ------------------------- Compression ---------------------- #

0 commit comments

Comments
 (0)