Instantiate NeRF's MLP model.
(args)
| 268 | self.alpha_linear.bias.data = torch.from_numpy(np.transpose(weights[idx_alpha_linear+1])) |
| 269 | |
| 270 | def create_nerf(args): |
| 271 | """Instantiate NeRF's MLP model. |
| 272 | """ |
| 273 | if args.reduce_embedding==2: # use DNeRF embedding |
| 274 | embed_fn, input_ch, embedder_obj = get_embedder(args.multires, args.i_embed, args.reduce_embedding, args.epochToMaxFreq) # input_ch.shape=63 |
| 275 | else: |
| 276 | embed_fn, input_ch, _ = get_embedder(args.multires, args.i_embed, args.reduce_embedding) # input_ch.shape=63 |
| 277 | |
| 278 | input_ch_views = 0 |
| 279 | embeddirs_fn = None |
| 280 | if args.use_viewdirs: |
| 281 | if args.reduce_embedding==2: # use DNeRF embedding |
| 282 | if args.no_DNeRF_viewdir: # no DNeRF embedding for viewdir |
| 283 | embeddirs_fn, input_ch_views, _ = get_embedder(args.multires_views, args.i_embed) |
| 284 | else: |
| 285 | embeddirs_fn, input_ch_views, embedddirs_obj = get_embedder(args.multires_views, args.i_embed, args.reduce_embedding, args.epochToMaxFreq) |
| 286 | else: |
| 287 | embeddirs_fn, input_ch_views, _ = get_embedder(args.multires_views, args.i_embed, args.reduce_embedding) # input_ch_views.shape=27 |
| 288 | output_ch = 5 if args.N_importance > 0 else 4 |
| 289 | skips = [4] |
| 290 | model = NeRF(D=args.netdepth, W=args.netwidth, input_ch=input_ch, output_ch=output_ch, skips=skips, |
| 291 | input_ch_views=input_ch_views, use_viewdirs=args.use_viewdirs) |
| 292 | device = torch.device("cuda") |
| 293 | if args.multi_gpu: |
| 294 | model = torch.nn.DataParallel(model).to(device) |
| 295 | else: |
| 296 | model = model.to(device) |
| 297 | grad_vars = list(model.parameters()) |
| 298 | |
| 299 | model_fine = None |
| 300 | if args.N_importance > 0: |
| 301 | model_fine = NeRF(D=args.netdepth_fine, W=args.netwidth_fine, input_ch=input_ch, output_ch=output_ch, skips=skips, |
| 302 | input_ch_views=input_ch_views, use_viewdirs=args.use_viewdirs) |
| 303 | if args.multi_gpu: |
| 304 | model_fine = torch.nn.DataParallel(model_fine).to(device) |
| 305 | else: |
| 306 | model_fine = model_fine.to(device) |
| 307 | grad_vars += list(model_fine.parameters()) |
| 308 | |
| 309 | if args.reduce_embedding==2: # use DNeRF embedding |
| 310 | network_query_fn = lambda inputs, viewdirs, network_fn, epoch: run_network_DNeRF(inputs, viewdirs, network_fn, |
| 311 | embed_fn=embed_fn, |
| 312 | embeddirs_fn=embeddirs_fn, |
| 313 | netchunk=args.netchunk, |
| 314 | epoch=epoch, no_DNeRF_viewdir=args.no_DNeRF_viewdir) |
| 315 | else: |
| 316 | network_query_fn = lambda inputs, viewdirs, network_fn : run_network(inputs, viewdirs, network_fn, |
| 317 | embed_fn=embed_fn, |
| 318 | embeddirs_fn=embeddirs_fn, |
| 319 | netchunk=args.netchunk) |
| 320 | |
| 321 | # Create optimizer |
| 322 | optimizer = torch.optim.Adam(params=grad_vars, lr=args.lrate, betas=(0.9, 0.999)) |
| 323 | |
| 324 | start = 0 |
| 325 | basedir = args.basedir |
| 326 | expname = args.expname |
| 327 |
no test coverage detected