Skip to content

Commit 6aed776

Browse files
sandeepnmenonsandeepnmenon
committed
test function for rasterszaton tests
1 parent b0dfe34 commit 6aed776

File tree

1 file changed

+167
-0
lines changed

1 file changed

+167
-0
lines changed

‎rasterization_tests.py

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
import math
2+
import time
3+
4+
import torch
5+
6+
from diff_gaussian_rasterization import (
7+
GaussianRasterizationSettings,
8+
GaussianRasterizer,
9+
)
10+
11+
12+
def get_cuda_args(strategy, mode="train"):
13+
cuda_args = {
14+
"mode": mode,
15+
"world_size": "1",
16+
"global_rank": "0",
17+
"local_rank": "0",
18+
"mp_world_size": "1",
19+
"mp_rank": "0",
20+
"log_folder": "./logs",
21+
"log_interval": "10",
22+
"iteration": "0",
23+
"zhx_debug": "False",
24+
"zhx_time": "False",
25+
"dist_global_strategy": "default",
26+
"avoid_pixel_all2all": False,
27+
"stats_collector": {},
28+
}
29+
return cuda_args
30+
31+
def test_gaussian_rasterizer_time():
32+
# Set up the input data
33+
num_gaussians = 10000
34+
means3D = torch.randn(num_gaussians, 3).cuda()
35+
scales = torch.randn(num_gaussians, 3).cuda()
36+
rotations = torch.randn(num_gaussians, 3, 3).cuda()
37+
shs = torch.randn(num_gaussians, 9).cuda()
38+
opacities = torch.randn(num_gaussians, 1).cuda()
39+
40+
# Set up the rasterization settings
41+
image_height = 512
42+
image_width = 512
43+
tanfovx = 1.0
44+
tanfovy = 1.0
45+
bg = torch.ones(3).cuda()
46+
scale_modifier = 1.0
47+
viewmatrix = torch.eye(4).cuda()
48+
projmatrix = torch.eye(4).cuda()
49+
sh_degree = 2
50+
campos = torch.zeros(3).cuda()
51+
prefiltered = False
52+
debug = False
53+
54+
# mode="train"
55+
# strategy=None
56+
# cuda_args = get_cuda_args(strategy, mode)
57+
58+
raster_settings = GaussianRasterizationSettings(
59+
image_height, image_width, tanfovx, tanfovy, bg,
60+
scale_modifier, viewmatrix, projmatrix, sh_degree,
61+
campos, prefiltered, debug
62+
)
63+
64+
# Create the GaussianRasterizer
65+
rasterizer = GaussianRasterizer(raster_settings)
66+
67+
# Measure the time for preprocess_gaussians
68+
start_time = time.time()
69+
means2D, rgb, conic_opacity, radii, depths = rasterizer.preprocess_gaussians(
70+
means3D, scales, rotations, shs, opacities
71+
)
72+
end_time = time.time()
73+
74+
preprocess_time = end_time - start_time
75+
print(f"Time taken by preprocess_gaussians: {preprocess_time:.4f} seconds")
76+
77+
78+
def test_batched_gaussian_rasterizer():
79+
# Set up the input data
80+
num_gaussians = 10000
81+
num_batches = 4
82+
means3D = torch.randn(num_gaussians, 3).cuda()
83+
scales = torch.randn(num_gaussians, 3).cuda()
84+
rotations = torch.randn(num_gaussians, 3, 3).cuda()
85+
shs = torch.randn(num_gaussians, 9).cuda()
86+
opacity = torch.randn(num_gaussians, 1).cuda()
87+
88+
# Set up the viewpoint cameras
89+
batched_viewpoint_cameras = []
90+
for _ in range(num_batches):
91+
viewpoint_camera = type('ViewpointCamera', (), {})
92+
viewpoint_camera.FoVx = math.radians(60)
93+
viewpoint_camera.FoVy = math.radians(60)
94+
viewpoint_camera.image_height = 512
95+
viewpoint_camera.image_width = 512
96+
viewpoint_camera.world_view_transform = torch.eye(4).cuda()
97+
viewpoint_camera.full_proj_transform = torch.eye(4).cuda()
98+
viewpoint_camera.camera_center = torch.zeros(3).cuda()
99+
batched_viewpoint_cameras.append(viewpoint_camera)
100+
101+
# Set up the strategies
102+
batched_strategies = [None] * num_batches
103+
104+
# Set up other parameters
105+
bg_color = torch.ones(3).cuda()
106+
scaling_modifier = 1.0
107+
pc = type('PC', (), {})
108+
pc.active_sh_degree = 2
109+
pipe = type('Pipe', (), {})
110+
pipe.debug = False
111+
mode = "train"
112+
113+
batched_rasterizers = []
114+
batched_cuda_args = []
115+
batched_screenspace_params = []
116+
batched_means2D = []
117+
batched_radii = []
118+
119+
for i, (viewpoint_camera, strategy) in enumerate(zip(batched_viewpoint_cameras, batched_strategies)):
120+
########## [START] Prepare CUDA Rasterization Settings ##########
121+
cuda_args = get_cuda_args(strategy, mode)
122+
batched_cuda_args.append(cuda_args)
123+
124+
# Set up rasterization configuration
125+
tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
126+
tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)
127+
raster_settings = GaussianRasterizationSettings(
128+
image_height=int(viewpoint_camera.image_height),
129+
image_width=int(viewpoint_camera.image_width),
130+
tanfovx=tanfovx,
131+
tanfovy=tanfovy,
132+
bg=bg_color,
133+
scale_modifier=scaling_modifier,
134+
viewmatrix=viewpoint_camera.world_view_transform,
135+
projmatrix=viewpoint_camera.full_proj_transform,
136+
sh_degree=pc.active_sh_degree,
137+
campos=viewpoint_camera.camera_center,
138+
prefiltered=False,
139+
debug=pipe.debug
140+
)
141+
rasterizer = GaussianRasterizer(raster_settings=raster_settings)
142+
########## [END] Prepare CUDA Rasterization Settings ##########
143+
144+
#[3DGS-wise preprocess]
145+
means2D, rgb, conic_opacity, radii, depths = rasterizer.preprocess_gaussians(
146+
means3D=means3D,
147+
scales=scales,
148+
rotations=rotations,
149+
shs=shs,
150+
opacities=opacity,
151+
cuda_args=cuda_args
152+
)
153+
154+
if mode == "train":
155+
means2D.retain_grad()
156+
157+
batched_means2D.append(means2D)
158+
screenspace_params = [means2D, rgb, conic_opacity, radii, depths]
159+
batched_rasterizers.append(rasterizer)
160+
batched_screenspace_params.append(screenspace_params)
161+
batched_radii.append(radii)
162+
163+
# Perform further operations with the batched results
164+
# ...
165+
166+
if __name__ == "__main__":
167+
test_gaussian_rasterizer_time()

0 commit comments

Comments
 (0)