| 310 | _CONFIG_SERIALIZERS = [_strategy_serializer] |
| 311 | |
| 312 | def create_splats_with_optimizers( |
| 313 | parser: Parser, |
| 314 | init_type: str = "sfm", |
| 315 | init_num_pts: int = 100_000, |
| 316 | init_extent: float = 3.0, |
| 317 | init_opacity: float = 0.1, |
| 318 | init_scale: float = 1.0, |
| 319 | scene_scale: float = 1.0, |
| 320 | sh_degree: int = 3, |
| 321 | sparse_grad: bool = False, |
| 322 | visible_adam: bool = False, |
| 323 | batch_size: int = 1, |
| 324 | feature_dim: Optional[int] = None, |
| 325 | device: str = "cuda", |
| 326 | world_rank: int = 0, |
| 327 | world_size: int = 1, |
| 328 | ) -> Tuple[torch.nn.ParameterDict, Dict[str, torch.optim.Optimizer]]: |
| 329 | if init_type == "sfm": |
| 330 | points = torch.from_numpy(parser.points).float() |
| 331 | rgbs = torch.from_numpy(parser.points_rgb / 255.0).float() |
| 332 | elif init_type == "random": |
| 333 | points = init_extent * scene_scale * (torch.rand((init_num_pts, 3)) * 2 - 1) |
| 334 | rgbs = torch.rand((init_num_pts, 3)) |
| 335 | else: |
| 336 | raise ValueError("Please specify a correct init_type: sfm or random") |
| 337 | |
| 338 | # Initialize the GS size to be the average dist of the 3 nearest neighbors |
| 339 | dist2_avg = (knn(points, 4)[:, 1:] ** 2).mean(dim=-1) # [N,] |
| 340 | dist_avg = torch.sqrt(dist2_avg) |
| 341 | scales = torch.log(dist_avg * init_scale).unsqueeze(-1).repeat(1, 3) # [N, 3] |
| 342 | |
| 343 | # Distribute the GSs to different ranks (also works for single rank) |
| 344 | points = points[world_rank::world_size] |
| 345 | rgbs = rgbs[world_rank::world_size] |
| 346 | scales = scales[world_rank::world_size] |
| 347 | |
| 348 | N = points.shape[0] |
| 349 | quats = torch.rand((N, 4)) # [N, 4] |
| 350 | opacities = torch.logit(torch.full((N,), init_opacity)) # [N,] |
| 351 | |
| 352 | params = [ |
| 353 | # name, value, lr |
| 354 | ("means", torch.nn.Parameter(points), 1.6e-4 * scene_scale), |
| 355 | ("scales", torch.nn.Parameter(scales), 5e-3), |
| 356 | ("quats", torch.nn.Parameter(quats), 1e-3), |
| 357 | ("opacities", torch.nn.Parameter(opacities), 5e-2), |
| 358 | ] |
| 359 | |
| 360 | if feature_dim is None: |
| 361 | # color is SH coefficients. |
| 362 | colors = torch.zeros((N, (sh_degree + 1) ** 2, 3)) # [N, K, 3] |
| 363 | colors[:, 0, :] = rgb_to_sh(rgbs) |
| 364 | params.append(("sh0", torch.nn.Parameter(colors[:, :1, :]), 2.5e-3)) |
| 365 | params.append(("shN", torch.nn.Parameter(colors[:, 1:, :]), 2.5e-3 / 20)) |
| 366 | else: |
| 367 | # features will be used for appearance and view-dependent shading |
| 368 | features = torch.rand(N, feature_dim) # [N, feature_dim] |
| 369 | params.append(("features", torch.nn.Parameter(features), 2.5e-3)) |