Skip to content

Commit

Permalink
Revert "rename to "proposal_resampling""
Browse files Browse the repository at this point in the history 
This reverts commit 3090d7b.
  • Loading branch information
@liruilong940607
liruilong940607 committed Nov 21, 2022
1 parent 3090d7b commit 55acb15
Show file tree
Hide file tree
Showing 11 changed files with 135 additions and 551 deletions.
60 changes: 20 additions & 40 deletions examples/radiance_fields/ngp.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
Copyright (c) 2022 Ruilong Li, UC Berkeley.
"""

import math
from typing import Callable, List, Union

import torch
Expand Down Expand Up @@ -74,11 +73,8 @@ def __init__(
use_viewdirs: bool = True,
density_activation: Callable = lambda x: trunc_exp(x - 1),
unbounded: bool = False,
hidden_dim: int = 64,
geo_feat_dim: int = 15,
n_levels: int = 16,
max_res: int = 1024,
base_res: int = 16,
log2_hashmap_size: int = 19,
) -> None:
super().__init__()
Expand All @@ -91,9 +87,7 @@ def __init__(
self.unbounded = unbounded

self.geo_feat_dim = geo_feat_dim
per_level_scale = math.exp(
(math.log(max_res) - math.log(base_res)) / (n_levels - 1)
)
per_level_scale = 1.4472692012786865

if self.use_viewdirs:
self.direction_encoding = tcnn.Encoding(
Expand All @@ -111,39 +105,25 @@ def __init__(
},
)

if hidden_dim > 0:
self.mlp_base = tcnn.NetworkWithInputEncoding(
n_input_dims=num_dim,
n_output_dims=1 + self.geo_feat_dim,
encoding_config={
"otype": "HashGrid",
"n_levels": n_levels,
"n_features_per_level": 2,
"log2_hashmap_size": log2_hashmap_size,
"base_resolution": base_res,
"per_level_scale": per_level_scale,
},
network_config={
"otype": "FullyFusedMLP",
"activation": "ReLU",
"output_activation": "None",
"n_neurons": hidden_dim,
"n_hidden_layers": 1,
},
)
else:
self.mlp_base = tcnn.Encoding(
n_input_dims=num_dim,
encoding_config={
"otype": "HashGrid",
"n_levels": 1,
"n_features_per_level": 1,
"log2_hashmap_size": 21,
"base_resolution": 128,
"per_level_scale": 1.0,
},
)

self.mlp_base = tcnn.NetworkWithInputEncoding(
n_input_dims=num_dim,
n_output_dims=1 + self.geo_feat_dim,
encoding_config={
"otype": "HashGrid",
"n_levels": n_levels,
"n_features_per_level": 2,
"log2_hashmap_size": log2_hashmap_size,
"base_resolution": 16,
"per_level_scale": per_level_scale,
},
network_config={
"otype": "FullyFusedMLP",
"activation": "ReLU",
"output_activation": "None",
"n_neurons": 64,
"n_hidden_layers": 1,
},
)
if self.geo_feat_dim > 0:
self.mlp_head = tcnn.Network(
n_input_dims=(
Expand Down
2 changes: 1 addition & 1 deletion nerfacc/cdf.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ def ray_resampling(
resampled_t_starts,
resampled_t_ends,
) = _C.ray_resampling(
packed_info.contiguous().int(),
packed_info.contiguous(),
t_starts.contiguous(),
t_ends.contiguous(),
weights.contiguous(),
Expand Down
2 changes: 0 additions & 2 deletions nerfacc/cuda/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,7 @@ def call_cuda(*args, **kwargs):

ray_aabb_intersect = _make_lazy_cuda_func("ray_aabb_intersect")
ray_marching = _make_lazy_cuda_func("ray_marching")
ray_marching_with_grid = _make_lazy_cuda_func("ray_marching_with_grid")
ray_resampling = _make_lazy_cuda_func("ray_resampling")
ray_pdf_query = _make_lazy_cuda_func("ray_pdf_query")

is_cub_available = _make_lazy_cuda_func("is_cub_available")
transmittance_from_sigma_forward_cub = _make_lazy_cuda_func(
Expand Down
183 changes: 21 additions & 162 deletions nerfacc/cuda/csrc/cdf.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,97 +4,13 @@

#include "include/helpers_cuda.h"

template <typename scalar_t>
__global__ void pdf_query_kernel(
const uint32_t n_rays,
// query
const int *packed_info, // input ray & point indices.
const scalar_t *starts, // input start t
const scalar_t *ends, // input end t
const scalar_t *pdfs, // pdf to be queried
// resample
const int *resample_packed_info, // input ray & point indices.
const scalar_t *resample_starts, // input start t, sorted
const scalar_t *resample_ends, // input end t, sorted
// output
scalar_t *resample_pdfs) // should be zero-initialized
{
CUDA_GET_THREAD_ID(i, n_rays);

// locate
const int base = packed_info[i * 2 + 0]; // point idx start.
const int steps = packed_info[i * 2 + 1]; // point idx shift.
const int resample_base = resample_packed_info[i * 2 + 0]; // point idx start.
const int resample_steps = resample_packed_info[i * 2 + 1]; // point idx shift.

if (resample_steps == 0) // nothing to query
return;

if (steps == 0) // nothing to be queried: set pdfs to 0
return;

starts += base;
ends += base;
pdfs += base;
resample_starts += resample_base;
resample_ends += resample_base;
resample_pdfs += resample_base;

// which interval is resample_start (t0) located
int t0_id = -1;
scalar_t t0_start = 0.0f, t0_end = starts[0];
scalar_t cdf0_start = 0.0f, cdf0_end = 0.0f;
// which interval is resample_end (t1) located
int t1_id = -1;
scalar_t t1_start = 0.0f, t1_end = starts[0];
scalar_t cdf1_start = 0.0f, cdf1_end = 0.0f;
// go!
for (int j = 0; j < resample_steps; ++j)
{
scalar_t t0 = resample_starts[j];
while(t0 > t0_end & t0_id < steps - 1) {
t0_id++;
t0_start = starts[t0_id];
t0_end = ends[t0_id];
cdf0_start = cdf0_end;
cdf0_end += pdfs[t0_id];
}
if (t0 > t0_end) {
resample_pdfs[j] = 0.0f;
continue;
}
scalar_t pct0 = 0.0f; // max(t0 - t0_start, 0.0f) / max(t0_end - t0_start, 1e-10f);
scalar_t resample_cdf_start = cdf0_start + pct0 * (cdf0_end - cdf0_start);

scalar_t t1 = resample_ends[j];
while(t1 > t1_end & t1_id < steps - 1) {
t1_id++;
t1_start = starts[t1_id];
t1_end = ends[t1_id];
cdf1_start = cdf1_end;
cdf1_end += pdfs[t1_id];
}
if (t1 > t1_end) {
resample_pdfs[j] = cdf1_end - resample_cdf_start;
continue;
}
scalar_t pct1 = 1.0f; // max(t1 - t1_start, 0.0f) / max(t1_end - t1_start, 1e-10f);
scalar_t resample_cdf_end = cdf1_start + pct1 * (cdf1_end - cdf1_start);

// compute pdf of [t0, t1]
resample_pdfs[j] = resample_cdf_end - resample_cdf_start;
}

return;
}

template <typename scalar_t>
__global__ void cdf_resampling_kernel(
const uint32_t n_rays,
const int *packed_info, // input ray & point indices.
const scalar_t *starts, // input start t
const scalar_t *ends, // input end t
const scalar_t *w, // transmittance weights
const scalar_t *weights, // transmittance weights
const int *resample_packed_info,
scalar_t *resample_starts,
scalar_t *resample_ends)
Expand All @@ -111,59 +27,52 @@ __global__ void cdf_resampling_kernel(

starts += base;
ends += base;
w += base;
weights += base;
resample_starts += resample_base;
resample_ends += resample_base;

// normalize weights **per ray**
scalar_t w_sum = 0.0f;
scalar_t weights_sum = 0.0f;
for (int j = 0; j < steps; j++)
w_sum += w[j];
// scalar_t padding = fmaxf(1e-10f - weights_sum, 0.0f);
// scalar_t padding_step = padding / steps;
// weights_sum += padding;
weights_sum += weights[j];
scalar_t padding = fmaxf(1e-5f - weights_sum, 0.0f);
scalar_t padding_step = padding / steps;
weights_sum += padding;

int num_endpoints = resample_steps + 1;
scalar_t cdf_pad = 1.0f / (2 * num_endpoints);
scalar_t cdf_step_size = (1.0f - 2 * cdf_pad) / resample_steps;
int num_bins = resample_steps + 1;
scalar_t cdf_step_size = (1.0f - 1.0 / num_bins) / resample_steps;

int idx = 0, j = 0;
scalar_t cdf_prev = 0.0f, cdf_next = w[idx] / w_sum;
scalar_t cdf_u = cdf_pad;
while (j < num_endpoints)
scalar_t cdf_prev = 0.0f, cdf_next = (weights[idx] + padding_step) / weights_sum;
scalar_t cdf_u = 1.0 / (2 * num_bins);
while (j < num_bins)
{
if (cdf_u < cdf_next)
{
// printf("cdf_u: %f, cdf_next: %f\n", cdf_u, cdf_next);
// resample in this interval
scalar_t scaling = (ends[idx] - starts[idx]) / (cdf_next - cdf_prev);
scalar_t t = (cdf_u - cdf_prev) * scaling + starts[idx];
// if (j == 100) {
// printf(
// "cdf_u: %.10f, cdf_next: %.10f, cdf_prev: %.10f, scaling: %.10f, t: %.10f, starts[idx]: %.10f, ends[idx]: %.10f\n",
// cdf_u, cdf_next, cdf_prev, scaling, t, starts[idx], ends[idx]);
// }
if (j < num_endpoints - 1)
if (j < num_bins - 1)
resample_starts[j] = t;
if (j > 0)
resample_ends[j - 1] = t;
// going further to next resample
// cdf_u += cdf_step_size;
cdf_u += cdf_step_size;
j += 1;
cdf_u = j * cdf_step_size + cdf_pad;
}
else
{
// going to next interval
idx += 1;
cdf_prev = cdf_next;
cdf_next += w[idx] / w_sum;
cdf_next += (weights[idx] + padding_step) / weights_sum;
}
}
// if (j != num_endpoints)
// {
// printf("Error: %d %d %f\n", j, num_endpoints, weights_sum);
// }
if (j != num_bins)
{
printf("Error: %d %d %f\n", j, num_bins, weights_sum);
}
return;
}

Expand Down Expand Up @@ -256,7 +165,7 @@ std::vector<torch::Tensor> ray_resampling(
TORCH_CHECK(packed_info.ndimension() == 2 & packed_info.size(1) == 2);
TORCH_CHECK(starts.ndimension() == 2 & starts.size(1) == 1);
TORCH_CHECK(ends.ndimension() == 2 & ends.size(1) == 1);
TORCH_CHECK(weights.ndimension() == 2 & weights.size(1) == 1);
TORCH_CHECK(weights.ndimension() == 1);

const uint32_t n_rays = packed_info.size(0);
const uint32_t n_samples = weights.size(0);
Expand All @@ -265,8 +174,7 @@ std::vector<torch::Tensor> ray_resampling(
const int blocks = CUDA_N_BLOCKS_NEEDED(n_rays, threads);

torch::Tensor num_steps = torch::split(packed_info, 1, 1)[1];
// torch::Tensor resample_num_steps = (num_steps > 0).to(num_steps.options()) * steps;
torch::Tensor resample_num_steps = torch::clamp(num_steps, 0, steps);
torch::Tensor resample_num_steps = (num_steps > 0).to(num_steps.options()) * steps;
torch::Tensor resample_cum_steps = resample_num_steps.cumsum(0, torch::kInt32);
torch::Tensor resample_packed_info = torch::cat(
{resample_cum_steps - resample_num_steps, resample_num_steps}, 1);
Expand All @@ -293,52 +201,3 @@ std::vector<torch::Tensor> ray_resampling(

return {resample_packed_info, resample_starts, resample_ends};
}

torch::Tensor ray_pdf_query(
torch::Tensor packed_info,
torch::Tensor starts,
torch::Tensor ends,
torch::Tensor pdfs,
torch::Tensor resample_packed_info,
torch::Tensor resample_starts,
torch::Tensor resample_ends)
{
DEVICE_GUARD(packed_info);

CHECK_INPUT(packed_info);
CHECK_INPUT(starts);
CHECK_INPUT(ends);
CHECK_INPUT(pdfs);

TORCH_CHECK(packed_info.ndimension() == 2 & packed_info.size(1) == 2);
TORCH_CHECK(starts.ndimension() == 2 & starts.size(1) == 1);
TORCH_CHECK(ends.ndimension() == 2 & ends.size(1) == 1);
TORCH_CHECK(pdfs.ndimension() == 2 & pdfs.size(1) == 1);

const uint32_t n_rays = packed_info.size(0);
const uint32_t n_resamples = resample_starts.size(0);

const int threads = 256;
const int blocks = CUDA_N_BLOCKS_NEEDED(n_rays, threads);

torch::Tensor resample_pdfs = torch::zeros({n_resamples, 1}, pdfs.options());

AT_DISPATCH_FLOATING_TYPES_AND_HALF(
pdfs.scalar_type(),
"pdf_query",
([&]
{ pdf_query_kernel<scalar_t><<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
n_rays,
// inputs
packed_info.data_ptr<int>(),
starts.data_ptr<scalar_t>(),
ends.data_ptr<scalar_t>(),
pdfs.data_ptr<scalar_t>(),
resample_packed_info.data_ptr<int>(),
resample_starts.data_ptr<scalar_t>(),
resample_ends.data_ptr<scalar_t>(),
// outputs
resample_pdfs.data_ptr<scalar_t>()); }));

return resample_pdfs;
}
18 changes: 0 additions & 18 deletions nerfacc/cuda/csrc/pybind.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,13 +13,6 @@ std::vector<torch::Tensor> ray_aabb_intersect(
const torch::Tensor aabb);

std::vector<torch::Tensor> ray_marching(
// rays
const torch::Tensor t_min,
const torch::Tensor t_max,
// sampling
const float step_size,
const float cone_angle);
std::vector<torch::Tensor> ray_marching_with_grid(
// rays
const torch::Tensor rays_o,
const torch::Tensor rays_d,
Expand Down Expand Up @@ -65,15 +58,6 @@ std::vector<torch::Tensor> ray_resampling(
torch::Tensor weights,
const int steps);

torch::Tensor ray_pdf_query(
torch::Tensor packed_info,
torch::Tensor starts,
torch::Tensor ends,
torch::Tensor pdfs,
torch::Tensor resample_packed_info,
torch::Tensor resample_starts,
torch::Tensor resample_ends);

torch::Tensor unpack_data(
torch::Tensor packed_info,
torch::Tensor data,
Expand Down Expand Up @@ -160,9 +144,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
// marching
m.def("ray_aabb_intersect", &ray_aabb_intersect);
m.def("ray_marching", &ray_marching);
m.def("ray_marching_with_grid", &ray_marching_with_grid);
m.def("ray_resampling", &ray_resampling);
m.def("ray_pdf_query", &ray_pdf_query);

// rendering
m.def("is_cub_available", is_cub_available);
Expand Down
Loading

0 comments on commit 55acb15

Please sign in to comment.