(args, train_dl, H, W, focal, N_rand, optimizer, loss_func, global_step, render_kwargs_train)
| 30 | args = parser.parse_args() |
| 31 | |
| 32 | def train_on_epoch_nerfw(args, train_dl, H, W, focal, N_rand, optimizer, loss_func, global_step, render_kwargs_train): |
| 33 | for batch_idx, (target, pose, img_idx) in enumerate(train_dl): |
| 34 | target = target[0].permute(1,2,0).to(device) |
| 35 | pose = pose.reshape(3,4).to(device) # reshape to 3x4 rot matrix |
| 36 | img_idx = img_idx.to(device) |
| 37 | |
| 38 | torch.set_default_tensor_type('torch.cuda.FloatTensor') |
| 39 | if N_rand is not None: |
| 40 | rays_o, rays_d = get_rays(H, W, focal, torch.Tensor(pose)) # (H, W, 3), (H, W, 3) |
| 41 | coords = torch.stack(torch.meshgrid(torch.linspace(0, H-1, H), torch.linspace(0, W-1, W), indexing='ij'), -1) # (H, W, 2) |
| 42 | coords = torch.reshape(coords, [-1,2]) # (H * W, 2) |
| 43 | select_inds = np.random.choice(coords.shape[0], size=[N_rand], replace=False) # (N_rand,) |
| 44 | select_coords = coords[select_inds].long() # (N_rand, 2) |
| 45 | rays_o = rays_o[select_coords[:, 0], select_coords[:, 1]] # (N_rand, 3) |
| 46 | rays_d = rays_d[select_coords[:, 0], select_coords[:, 1]] # (N_rand, 3) |
| 47 | batch_rays = torch.stack([rays_o, rays_d], 0) |
| 48 | target_s = target[select_coords[:, 0], select_coords[:, 1]] # (N_rand, 3) |
| 49 | |
| 50 | # ##### Core optimization loop ##### |
| 51 | rgb, disp, acc, extras = render(H, W, focal, chunk=args.chunk, rays=batch_rays, retraw=True, img_idx=img_idx, **render_kwargs_train) |
| 52 | optimizer.zero_grad() |
| 53 | |
| 54 | # compute loss |
| 55 | results = {} |
| 56 | results['rgb_fine'] = rgb |
| 57 | results['rgb_coarse'] = extras['rgb0'] |
| 58 | results['beta'] = extras['beta'] |
| 59 | results['transient_sigmas'] = extras['transient_sigmas'] |
| 60 | |
| 61 | loss_d = loss_func(results, target_s) |
| 62 | loss = sum(l for l in loss_d.values()) |
| 63 | |
| 64 | with torch.no_grad(): |
| 65 | img_loss = img2mse(rgb, target_s) |
| 66 | psnr = mse2psnr(img_loss) |
| 67 | loss.backward() |
| 68 | optimizer.step() |
| 69 | |
| 70 | # NOTE: IMPORTANT! |
| 71 | ### update learning rate ### |
| 72 | decay_rate = 0.1 |
| 73 | decay_steps = args.lrate_decay * 1000 |
| 74 | new_lrate = args.lrate * (decay_rate ** (global_step / decay_steps)) |
| 75 | for param_group in optimizer.param_groups: |
| 76 | param_group['lr'] = new_lrate |
| 77 | ################################ |
| 78 | |
| 79 | torch.set_default_tensor_type('torch.FloatTensor') |
| 80 | return loss, psnr |
| 81 | |
| 82 | def train_nerf(args, train_dl, val_dl, hwf, i_split, near, far, render_poses=None, render_img=None): |
| 83 |
no test coverage detected