trainval.py 7.95 KB
Newer Older
Ruilong Li's avatar
Ruilong Li committed
1
2
3
4
5
6
7
import math
import time

import numpy as np
import torch
import torch.nn.functional as F
import tqdm
Ruilong Li's avatar
Ruilong Li committed
8
from datasets.nerf_synthetic import SubjectLoader, namedtuple_map
Ruilong Li's avatar
Ruilong Li committed
9
10
11
12
from radiance_fields.ngp import NGPradianceField

from nerfacc import OccupancyField, volumetric_rendering

Ruilong Li's avatar
Ruilong Li committed
13
TARGET_SAMPLE_BATCH_SIZE = 1 << 16
Ruilong Li's avatar
Ruilong Li committed
14

Ruilong Li's avatar
Ruilong Li committed
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
# import tqdm

# device = "cuda:0"
# radiance_field = NGPradianceField(aabb=[0, 0, 0, 1, 1, 1]).to(device)
# positions = torch.rand((TARGET_SAMPLE_BATCH_SIZE, 3), device=device)
# directions = torch.rand(positions.shape, device=device)
# optimizer = torch.optim.Adam(
#     radiance_field.parameters(),
#     lr=1e-10,
#     # betas=(0.9, 0.99),
#     eps=1e-15,
#     # weight_decay=1e-6,
# )
# for _ in tqdm.tqdm(range(1000)):
#     rgbs, sigmas = radiance_field(positions, directions)
#     loss = rgbs.mean()
#     optimizer.zero_grad()
#     loss.backward()
#     optimizer.step()
# exit()

Ruilong Li's avatar
Ruilong Li committed
36
37

def render_image(radiance_field, rays, render_bkgd, render_step_size):
Ruilong Li's avatar
Ruilong Li committed
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
    """Render the pixels of an image.

    Args:
      radiance_field: the radiance field of nerf.
      rays: a `Rays` namedtuple, the rays to be rendered.

    Returns:
      rgb: torch.tensor, rendered color image.
      depth: torch.tensor, rendered depth image.
      acc: torch.tensor, rendered accumulated weights per pixel.
    """
    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
    results = []
57
    chunk = torch.iinfo(torch.int32).max if radiance_field.training else 81920
Ruilong Li's avatar
Ruilong Li committed
58
59
    for i in range(0, num_rays, chunk):
        chunk_rays = namedtuple_map(lambda r: r[i : i + chunk], rays)
Ruilong Li's avatar
Ruilong Li committed
60
        chunk_results = volumetric_rendering(
Ruilong Li's avatar
Ruilong Li committed
61
62
63
64
65
66
67
            query_fn=radiance_field.forward,  # {x, dir} -> {rgb, density}
            rays_o=chunk_rays.origins,
            rays_d=chunk_rays.viewdirs,
            scene_aabb=occ_field.aabb,
            scene_occ_binary=occ_field.occ_grid_binary,
            scene_resolution=occ_field.resolution,
            render_bkgd=render_bkgd,
Ruilong Li's avatar
Ruilong Li committed
68
            render_step_size=render_step_size,
Ruilong Li's avatar
Ruilong Li committed
69
        )
Ruilong Li's avatar
Ruilong Li committed
70
        results.append(chunk_results)
Ruilong Li's avatar
Ruilong Li committed
71
    rgb, depth, acc, counter, compact_counter = [
Ruilong Li's avatar
Ruilong Li committed
72
73
        torch.cat(r, dim=0) if isinstance(r[0], torch.Tensor) else r
        for r in zip(*results)
Ruilong Li's avatar
Ruilong Li committed
74
    ]
Ruilong Li's avatar
Ruilong Li committed
75
76
77
78
    return (
        rgb.view((*rays_shape[:-1], -1)),
        depth.view((*rays_shape[:-1], -1)),
        acc.view((*rays_shape[:-1], -1)),
Ruilong Li's avatar
Ruilong Li committed
79
80
        sum(counter),
        sum(compact_counter),
Ruilong Li's avatar
Ruilong Li committed
81
82
83
84
    )


if __name__ == "__main__":
85
    torch.manual_seed(42)
Ruilong Li's avatar
Ruilong Li committed
86
87

    device = "cuda:0"
Ruilong Li's avatar
Ruilong Li committed
88
    scene = "lego"
Ruilong Li's avatar
Ruilong Li committed
89
90
91

    # setup dataset
    train_dataset = SubjectLoader(
Ruilong Li's avatar
Ruilong Li committed
92
        subject_id=scene,
Ruilong Li's avatar
Ruilong Li committed
93
        root_fp="/home/ruilongli/data/nerf_synthetic/",
Ruilong Li's avatar
wtf  
Ruilong Li committed
94
        split="trainval",
Ruilong Li's avatar
Ruilong Li committed
95
        num_rays=1024,
Ruilong Li's avatar
Ruilong Li committed
96
    )
Ruilong Li's avatar
Ruilong Li committed
97
98
99
100

    train_dataset.images = train_dataset.images.to(device)
    train_dataset.camtoworlds = train_dataset.camtoworlds.to(device)
    train_dataset.K = train_dataset.K.to(device)
Ruilong Li's avatar
Ruilong Li committed
101
102
    train_dataloader = torch.utils.data.DataLoader(
        train_dataset,
Ruilong Li's avatar
Ruilong Li committed
103
        num_workers=0,
Ruilong Li's avatar
Ruilong Li committed
104
        batch_size=None,
Ruilong Li's avatar
Ruilong Li committed
105
        # persistent_workers=True,
Ruilong Li's avatar
Ruilong Li committed
106
        shuffle=True,
Ruilong Li's avatar
Ruilong Li committed
107
    )
Ruilong Li's avatar
Ruilong Li committed
108

Ruilong Li's avatar
Ruilong Li committed
109
    test_dataset = SubjectLoader(
Ruilong Li's avatar
Ruilong Li committed
110
        subject_id=scene,
Ruilong Li's avatar
Ruilong Li committed
111
        root_fp="/home/ruilongli/data/nerf_synthetic/",
Ruilong Li's avatar
Ruilong Li committed
112
        split="test",
Ruilong Li's avatar
Ruilong Li committed
113
114
        num_rays=None,
    )
Ruilong Li's avatar
Ruilong Li committed
115
116
117
    test_dataset.images = test_dataset.images.to(device)
    test_dataset.camtoworlds = test_dataset.camtoworlds.to(device)
    test_dataset.K = test_dataset.K.to(device)
Ruilong Li's avatar
Ruilong Li committed
118
119
    test_dataloader = torch.utils.data.DataLoader(
        test_dataset,
Ruilong Li's avatar
Ruilong Li committed
120
        num_workers=0,
Ruilong Li's avatar
Ruilong Li committed
121
        batch_size=None,
Ruilong Li's avatar
Ruilong Li committed
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
    )

    # setup the scene bounding box.
    scene_aabb = torch.tensor([-1.5, -1.5, -1.5, 1.5, 1.5, 1.5])

    # setup the scene radiance field. Assume you have a NeRF model and
    # it has following functions:
    # - query_density(): {x} -> {density}
    # - forward(): {x, dirs} -> {rgb, density}
    radiance_field = NGPradianceField(aabb=scene_aabb).to(device)

    # setup some rendering settings
    render_n_samples = 1024
    render_step_size = (
        (scene_aabb[3:] - scene_aabb[:3]).max() * math.sqrt(3) / render_n_samples
Ruilong Li's avatar
Ruilong Li committed
137
    ).item()
Ruilong Li's avatar
Ruilong Li committed
138

Ruilong Li's avatar
Ruilong Li committed
139
140
141
142
143
144
145
    optimizer = torch.optim.Adam(
        radiance_field.parameters(),
        lr=1e-2,
        # betas=(0.9, 0.99),
        eps=1e-15,
        # weight_decay=1e-6,
    )
Ruilong Li's avatar
Ruilong Li committed
146
147
148
    scheduler = torch.optim.lr_scheduler.MultiStepLR(
        optimizer, milestones=[20000, 30000], gamma=0.1
    )
Ruilong Li's avatar
Ruilong Li committed
149
150
151
152
153
154
155
156
157
158
159

    # setup occupancy field with eval function
    def occ_eval_fn(x: torch.Tensor) -> torch.Tensor:
        """Evaluate occupancy given positions.

        Args:
            x: positions with shape (N, 3).
        Returns:
            occupancy values with shape (N, 1).
        """
        density_after_activation = radiance_field.query_density(x)
Ruilong Li's avatar
Ruilong Li committed
160
        # those two are similar when density is small.
161
        # occupancy = 1.0 - torch.exp(-density_after_activation * render_step_size)
Ruilong Li's avatar
Ruilong Li committed
162
163
164
165
166
167
168
169
170
171
        occupancy = density_after_activation * render_step_size
        return occupancy

    occ_field = OccupancyField(
        occ_eval_fn=occ_eval_fn, aabb=scene_aabb, resolution=128
    ).to(device)

    # training
    step = 0
    tic = time.time()
Ruilong Li's avatar
Ruilong Li committed
172
173
    data_time = 0
    tic_data = time.time()
Ruilong Li's avatar
wtf  
Ruilong Li committed
174

Ruilong Li's avatar
Ruilong Li committed
175
    for epoch in range(10000000):
Ruilong Li's avatar
Ruilong Li committed
176
177
        for i in range(len(train_dataset)):
            data = train_dataset[i]
Ruilong Li's avatar
Ruilong Li committed
178
            data_time += time.time() - tic_data
Ruilong Li's avatar
Ruilong Li committed
179
180

            # generate rays from data and the gt pixel color
Ruilong Li's avatar
Ruilong Li committed
181
182
183
184
185
            # rays = namedtuple_map(lambda x: x.to(device), data["rays"])
            # pixels = data["pixels"].to(device)
            render_bkgd = data["color_bkgd"]
            rays = data["rays"]
            pixels = data["pixels"]
Ruilong Li's avatar
Ruilong Li committed
186

Ruilong Li's avatar
Ruilong Li committed
187
188
            # update occupancy grid
            occ_field.every_n_step(step)
Ruilong Li's avatar
wtf  
Ruilong Li committed
189

Ruilong Li's avatar
Ruilong Li committed
190
            rgb, depth, acc, counter, compact_counter = render_image(
Ruilong Li's avatar
Ruilong Li committed
191
                radiance_field, rays, render_bkgd, render_step_size
Ruilong Li's avatar
readme  
Ruilong Li committed
192
            )
Ruilong Li's avatar
Ruilong Li committed
193
194
            num_rays = len(pixels)
            num_rays = int(
Ruilong Li's avatar
Ruilong Li committed
195
                num_rays * (TARGET_SAMPLE_BATCH_SIZE / float(compact_counter))
Ruilong Li's avatar
Ruilong Li committed
196
197
            )
            train_dataset.update_num_rays(num_rays)
Ruilong Li's avatar
Ruilong Li committed
198
            alive_ray_mask = acc.squeeze(-1) > 0
Ruilong Li's avatar
Ruilong Li committed
199

Ruilong Li's avatar
Ruilong Li committed
200
201
            # compute loss
            loss = F.mse_loss(rgb[alive_ray_mask], pixels[alive_ray_mask])
Ruilong Li's avatar
Ruilong Li committed
202

Ruilong Li's avatar
Ruilong Li committed
203
204
205
206
            optimizer.zero_grad()
            (loss * 128).backward()
            optimizer.step()
            scheduler.step()
Ruilong Li's avatar
Ruilong Li committed
207

Ruilong Li's avatar
Ruilong Li committed
208
            if step % 100 == 0:
Ruilong Li's avatar
Ruilong Li committed
209
210
                elapsed_time = time.time() - tic
                print(
Ruilong Li's avatar
Ruilong Li committed
211
                    f"elapsed_time={elapsed_time:.2f}s (data={data_time:.2f}s) | {step=} | "
Ruilong Li's avatar
Ruilong Li committed
212
213
                    f"loss={loss:.5f} | "
                    f"alive_ray_mask={alive_ray_mask.long().sum():d} | "
Ruilong Li's avatar
Ruilong Li committed
214
                    f"counter={counter:d} | compact_counter={compact_counter:d} | num_rays={len(pixels):d} "
Ruilong Li's avatar
Ruilong Li committed
215
216
                )

Ruilong Li's avatar
Ruilong Li committed
217
218
219
220
221
222
223
224
225
226
227
228
229
            # if time.time() - tic > 300:
            if step == 35_000:
                print("training stops")
                # evaluation
                radiance_field.eval()
                psnrs = []
                with torch.no_grad():
                    for data in tqdm.tqdm(test_dataloader):
                        # generate rays from data and the gt pixel color
                        rays = namedtuple_map(lambda x: x.to(device), data["rays"])
                        pixels = data["pixels"].to(device)
                        render_bkgd = data["color_bkgd"].to(device)
                        # rendering
Ruilong Li's avatar
Ruilong Li committed
230
                        rgb, depth, acc, _, _ = render_image(
Ruilong Li's avatar
Ruilong Li committed
231
232
233
234
235
236
237
238
                            radiance_field, rays, render_bkgd, render_step_size
                        )
                        mse = F.mse_loss(rgb, pixels)
                        psnr = -10.0 * torch.log(mse) / np.log(10.0)
                        psnrs.append(psnr.item())
                psnr_avg = sum(psnrs) / len(psnrs)
                print(f"evaluation: {psnr_avg=}")
                exit()
Ruilong Li's avatar
Ruilong Li committed
239
            tic_data = time.time()
Ruilong Li's avatar
Ruilong Li committed
240

Ruilong Li's avatar
Ruilong Li committed
241
            step += 1