Skip to content

GitLab

"git@developer.sourcefind.cn:OpenDAS/megatron-lm.git" did not exist on "74ba3797edb8ec41d2eceed0095dcaf3d6cedad8"
Unverified Commit 4d595ae4 authored by Ruilong Li(李瑞龙)'s avatar Ruilong Li(李瑞龙) Committed by GitHub
Browse files

rm build dir for first-time build (#137)

parent 903d3b55
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 556 deletions
examples/train_ngp_nerf_proposal.py deleted 100644 → 0
View file @ 903d3b55
"""
Copyright (c) 2022 Ruilong Li, UC Berkeley.
"""
import argparse
import math
import os
import random
import time
from typing import Optional
import imageio
import numpy as np
import torch
import torch.nn.functional as F
import tqdm
from datasets.utils import Rays, namedtuple_map
from radiance_fields.ngp import NGPradianceField
from utils import set_random_seed
from nerfacc import ContractionType, pack_info, ray_marching, rendering
from nerfacc.cuda import ray_pdf_query
def set_random_seed(seed):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
# @profile
def render_image(
# scene
radiance_field: torch.nn.Module,
proposal_nets: torch.nn.Module,
rays: Rays,
scene_aabb: torch.Tensor,
# rendering options
near_plane: Optional[float] = None,
far_plane: Optional[float] = None,
render_step_size: float = 1e-3,
render_bkgd: Optional[torch.Tensor] = None,
cone_angle: float = 0.0,
alpha_thre: float = 0.0,
proposal_nets_require_grads: bool = True,
# test options
test_chunk_size: int = 8192,
):
"""Render the pixels of an image."""
rays_shape = rays.origins.shape
if len(rays_shape) == 3:
height, width, _ = rays_shape
num_rays = height * width
rays = namedtuple_map(
lambda r: r.reshape([num_rays] + list(r.shape[2:])), rays
)
else:
num_rays, _ = rays_shape
def sigma_fn(t_starts, t_ends, ray_indices, net=None):
t_origins = chunk_rays.origins[ray_indices]
t_dirs = chunk_rays.viewdirs[ray_indices]
positions = t_origins + t_dirs * (t_starts + t_ends) / 2.0
if net is not None:
return net.query_density(positions)
else:
return radiance_field.query_density(positions)
def rgb_sigma_fn(t_starts, t_ends, ray_indices):
t_origins = chunk_rays.origins[ray_indices]
t_dirs = chunk_rays.viewdirs[ray_indices]
positions = t_origins + t_dirs * (t_starts + t_ends) / 2.0
return radiance_field(positions, t_dirs)
results = []
chunk = (
torch.iinfo(torch.int32).max
if radiance_field.training
else test_chunk_size
)
for i in range(0, num_rays, chunk):
chunk_rays = namedtuple_map(lambda r: r[i : i + chunk], rays)
ray_indices, t_starts, t_ends, proposal_sample_list = ray_marching(
chunk_rays.origins,
chunk_rays.viewdirs,
scene_aabb=scene_aabb,
grid=None,
# proposal density fns: {t_starts, t_ends, ray_indices} -> density
proposal_sigma_fns=[
lambda t_starts, t_ends, ray_indices: sigma_fn(
t_starts, t_ends, ray_indices, proposal_net
)
for proposal_net in proposal_nets
],
proposal_n_samples=[32],
proposal_require_grads=proposal_nets_require_grads,
sigma_fn=sigma_fn,
near_plane=near_plane,
far_plane=far_plane,
render_step_size=render_step_size,
stratified=radiance_field.training,
cone_angle=cone_angle,
alpha_thre=alpha_thre,
)
rgb, opacity, depth, weights = rendering(
t_starts,
t_ends,
ray_indices=ray_indices,
n_rays=len(chunk_rays.origins),
rgb_sigma_fn=rgb_sigma_fn,
render_bkgd=render_bkgd,
)
if radiance_field.training:
packed_info = pack_info(ray_indices, n_rays=len(chunk_rays.origins))
proposal_sample_list.append(
(packed_info, t_starts, t_ends, weights)
)
chunk_results = [rgb, opacity, depth, len(t_starts)]
results.append(chunk_results)
colors, opacities, depths, n_rendering_samples = [
torch.cat(r, dim=0) if isinstance(r[0], torch.Tensor) else r
for r in zip(*results)
]
return (
colors.view((*rays_shape[:-1], -1)),
opacities.view((*rays_shape[:-1], -1)),
depths.view((*rays_shape[:-1], -1)),
sum(n_rendering_samples),
proposal_sample_list if radiance_field.training else None,
)
if __name__ == "__main__":
device = "cuda:0"
set_random_seed(42)
parser = argparse.ArgumentParser()
parser.add_argument(
"--train_split",
type=str,
default="trainval",
choices=["train", "trainval"],
help="which train split to use",
)
parser.add_argument(
"--scene",
type=str,
default="lego",
choices=[
# nerf synthetic
"chair",
"drums",
"ficus",
"hotdog",
"lego",
"materials",
"mic",
"ship",
# mipnerf360 unbounded
"garden",
"bicycle",
"bonsai",
"counter",
"kitchen",
"room",
"stump",
],
help="which scene to use",
)
parser.add_argument(
"--aabb",
type=lambda s: [float(item) for item in s.split(",")],
default="-1.5,-1.5,-1.5,1.5,1.5,1.5",
help="delimited list input",
)
parser.add_argument(
"--test_chunk_size",
type=int,
default=1024,
)
parser.add_argument(
"--unbounded",
action="store_true",
help="whether to use unbounded rendering",
)
parser.add_argument(
"--auto_aabb",
action="store_true",
help="whether to automatically compute the aabb",
)
parser.add_argument("--cone_angle", type=float, default=0.0)
args = parser.parse_args()
render_n_samples = 256
# setup the dataset
train_dataset_kwargs = {}
test_dataset_kwargs = {}
if args.unbounded:
from datasets.nerf_360_v2 import SubjectLoader
data_root_fp = "/home/ruilongli/data/360_v2/"
target_sample_batch_size = 1 << 20
train_dataset_kwargs = {"color_bkgd_aug": "random", "factor": 4}
test_dataset_kwargs = {"factor": 4}
else:
from datasets.nerf_synthetic import SubjectLoader
data_root_fp = "/home/ruilongli/data/nerf_synthetic/"
target_sample_batch_size = 1 << 18
train_dataset = SubjectLoader(
subject_id=args.scene,
root_fp=data_root_fp,
split=args.train_split,
num_rays=target_sample_batch_size // 32,
**train_dataset_kwargs,
)
train_dataset.images = train_dataset.images.to(device)
train_dataset.camtoworlds = train_dataset.camtoworlds.to(device)
train_dataset.K = train_dataset.K.to(device)
test_dataset = SubjectLoader(
subject_id=args.scene,
root_fp=data_root_fp,
split="test",
num_rays=None,
**test_dataset_kwargs,
)
test_dataset.images = test_dataset.images.to(device)
test_dataset.camtoworlds = test_dataset.camtoworlds.to(device)
test_dataset.K = test_dataset.K.to(device)
if args.auto_aabb:
camera_locs = torch.cat(
[train_dataset.camtoworlds, test_dataset.camtoworlds]
)[:, :3, -1]
args.aabb = torch.cat(
[camera_locs.min(dim=0).values, camera_locs.max(dim=0).values]
).tolist()
print("Using auto aabb", args.aabb)
# setup the scene bounding box.
if args.unbounded:
print("Using unbounded rendering")
contraction_type = ContractionType.UN_BOUNDED_SPHERE
# contraction_type = ContractionType.UN_BOUNDED_TANH
scene_aabb = None
near_plane = 0.2
far_plane = 1e4
render_step_size = 1e-2
alpha_thre = 1e-2
else:
contraction_type = ContractionType.AABB
scene_aabb = torch.tensor(args.aabb, dtype=torch.float32, device=device)
near_plane = None
far_plane = None
render_step_size = (
(scene_aabb[3:] - scene_aabb[:3]).max()
* math.sqrt(3)
/ render_n_samples
).item()
alpha_thre = 1e-2
proposal_nets = torch.nn.ModuleList(
[
NGPradianceField(
aabb=args.aabb,
use_viewdirs=False,
hidden_dim=0,
geo_feat_dim=0,
),
# NGPradianceField(
# aabb=args.aabb,
# use_viewdirs=False,
# hidden_dim=16,
# max_res=64,
# geo_feat_dim=0,
# n_levels=4,
# log2_hashmap_size=19,
# ),
# NGPradianceField(
# aabb=args.aabb,
# use_viewdirs=False,
# hidden_dim=16,
# max_res=256,
# geo_feat_dim=0,
# n_levels=5,
# log2_hashmap_size=17,
# ),
]
).to(device)
# setup the radiance field we want to train.
max_steps = 20000
grad_scaler = torch.cuda.amp.GradScaler(2**10)
radiance_field = NGPradianceField(
aabb=args.aabb,
unbounded=args.unbounded,
).to(device)
optimizer = torch.optim.Adam(
list(radiance_field.parameters()) + list(proposal_nets.parameters()),
lr=1e-2,
eps=1e-15,
)
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
milestones=[max_steps // 2, max_steps * 3 // 4, max_steps * 9 // 10],
gamma=0.33,
)
# training
step = 0
tic = time.time()
for epoch in range(10000000):
for i in range(len(train_dataset)):
radiance_field.train()
proposal_nets.train()
# @profile
def _train():
data = train_dataset[i]
render_bkgd = data["color_bkgd"]
rays = data["rays"]
pixels = data["pixels"]
# render
(
rgb,
acc,
depth,
n_rendering_samples,
proposal_sample_list,
) = render_image(
radiance_field,
proposal_nets,
rays,
scene_aabb,
# rendering options
near_plane=near_plane,
far_plane=far_plane,
render_step_size=render_step_size,
render_bkgd=render_bkgd,
cone_angle=args.cone_angle,
alpha_thre=min(alpha_thre, alpha_thre * step / 1000),
proposal_nets_require_grads=(step < 100 or step % 16 == 0),
)
# if n_rendering_samples == 0:
# continue
# dynamic batch size for rays to keep sample batch size constant.
num_rays = len(pixels)
num_rays = int(
num_rays
* (target_sample_batch_size / float(n_rendering_samples))
)
train_dataset.update_num_rays(num_rays)
alive_ray_mask = acc.squeeze(-1) > 0
# compute loss
loss = F.smooth_l1_loss(
rgb[alive_ray_mask], pixels[alive_ray_mask]
)
(
packed_info,
t_starts,
t_ends,
weights,
) = proposal_sample_list[-1]
loss_interval = 0.0
for (
proposal_packed_info,
proposal_t_starts,
proposal_t_ends,
proposal_weights,
) in proposal_sample_list[:-1]:
proposal_weights_gt = ray_pdf_query(
packed_info,
t_starts,
t_ends,
weights.detach(),
proposal_packed_info,
proposal_t_starts,
proposal_t_ends,
).detach()
loss_interval = (
torch.clamp(
proposal_weights_gt - proposal_weights, min=0
)
) ** 2 / (proposal_weights + torch.finfo(torch.float32).eps)
loss_interval = loss_interval.mean()
loss += loss_interval * 1.0
optimizer.zero_grad()
# do not unscale it because we are using Adam.
grad_scaler.scale(loss).backward()
optimizer.step()
scheduler.step()
if step % 100 == 0:
elapsed_time = time.time() - tic
loss = F.mse_loss(
rgb[alive_ray_mask], pixels[alive_ray_mask]
)
print(
f"elapsed_time={elapsed_time:.2f}s | step={step} | "
f"loss={loss:.5f} | loss_interval={loss_interval:.5f} "
f"alive_ray_mask={alive_ray_mask.long().sum():d} | "
f"n_rendering_samples={n_rendering_samples:d} | num_rays={len(pixels):d} |"
)
_train()
if step >= 0 and step % 1000 == 0 and step > 0:
# evaluation
radiance_field.eval()
proposal_nets.eval()
psnrs = []
with torch.no_grad():
for i in tqdm.tqdm(range(len(test_dataset))):
data = test_dataset[i]
render_bkgd = data["color_bkgd"]
rays = data["rays"]
pixels = data["pixels"]
# rendering
rgb, acc, depth, _, _ = render_image(
radiance_field,
proposal_nets,
rays,
scene_aabb,
# rendering options
near_plane=near_plane,
far_plane=far_plane,
render_step_size=render_step_size,
render_bkgd=render_bkgd,
cone_angle=args.cone_angle,
alpha_thre=alpha_thre,
proposal_nets_require_grads=False,
# test options
test_chunk_size=args.test_chunk_size,
)
mse = F.mse_loss(rgb, pixels)
psnr = -10.0 * torch.log(mse) / np.log(10.0)
psnrs.append(psnr.item())
imageio.imwrite(
"acc_binary_test.png",
((acc > 0).float().cpu().numpy() * 255).astype(
np.uint8
),
)
imageio.imwrite(
"rgb_test.png",
(rgb.cpu().numpy() * 255).astype(np.uint8),
)
break
psnr_avg = sum(psnrs) / len(psnrs)
print(f"evaluation: psnr_avg={psnr_avg}")
train_dataset.training = True
if step == max_steps:
print("training stops")
exit()
step += 1
nerfacc/cuda/_backend.py
View file @ 4d595ae4
......@@ -75,6 +75,9 @@ if cuda_toolkit_available():
extra_include_paths=extra_include_paths,
)
else:
# Build from scratch. Remove the build directory just to be safe: pytorch jit might stuck
# if the build directory exists.
shutil.rmtree(build_dir)
with Console().status(
"[bold yellow]NerfAcc: Setting up CUDA (This may take a few minutes the first time)",
spinner="bouncingBall",
......
tests/test_pdf_query.py deleted 100644 → 0
View file @ 903d3b55
import torch
from torch import Tensor
from nerfacc import pack_info, ray_marching, ray_resampling
from nerfacc.cuda import ray_pdf_query
device = "cuda:0"
def outer(
t0_starts: Tensor,
t0_ends: Tensor,
t1_starts: Tensor,
t1_ends: Tensor,
y1: Tensor,
) -> Tensor:
cy1 = torch.cat(
[torch.zeros_like(y1[..., :1]), torch.cumsum(y1, dim=-1)], dim=-1
)
idx_lo = (
torch.searchsorted(
t1_starts.contiguous(), t0_starts.contiguous(), side="right"
)
- 1
)
idx_lo = torch.clamp(idx_lo, min=0, max=y1.shape[-1] - 1)
idx_hi = torch.searchsorted(
t1_ends.contiguous(), t0_ends.contiguous(), side="right"
)
idx_hi = torch.clamp(idx_hi, min=0, max=y1.shape[-1] - 1)
cy1_lo = torch.take_along_dim(cy1[..., :-1], idx_lo, dim=-1)
cy1_hi = torch.take_along_dim(cy1[..., 1:], idx_hi, dim=-1)
y0_outer = cy1_hi - cy1_lo
return y0_outer
def test_pdf_query():
n_rays = 1
rays_o = torch.rand((n_rays, 3), device=device)
rays_d = torch.randn((n_rays, 3), device=device)
rays_d = rays_d / rays_d.norm(dim=-1, keepdim=True)
ray_indices, t_starts, t_ends = ray_marching(
rays_o,
rays_d,
near_plane=0.1,
far_plane=1.0,
render_step_size=0.2,
)
packed_info = pack_info(ray_indices, n_rays)
weights = torch.rand((t_starts.shape[0], 1), device=device)
packed_info_new = packed_info
t_starts_new = t_starts - 0.3
t_ends_new = t_ends - 0.3
weights_new_ref = outer(
t_starts_new.reshape(n_rays, -1),
t_ends_new.reshape(n_rays, -1),
t_starts.reshape(n_rays, -1),
t_ends.reshape(n_rays, -1),
weights.reshape(n_rays, -1),
)
weights_new_ref = weights_new_ref.flatten()
weights_new = ray_pdf_query(
packed_info,
t_starts,
t_ends,
weights,
packed_info_new,
t_starts_new,
t_ends_new,
)
weights_new = weights_new.flatten()
# print(weights)
# print(weights_new_ref)
# print(weights_new)
if __name__ == "__main__":
test_pdf_query()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment