(self, batch, batch_idx)
| 193 | return out |
| 194 | |
| 195 | def training_step(self, batch, batch_idx): |
| 196 | lrm_generator_input, render_gt = self.prepare_batch_data(batch) |
| 197 | |
| 198 | render_out = self.forward(lrm_generator_input) |
| 199 | |
| 200 | loss, loss_dict = self.compute_loss(render_out, render_gt) |
| 201 | |
| 202 | self.log_dict(loss_dict, prog_bar=True, logger=True, on_step=True, on_epoch=True) |
| 203 | |
| 204 | if self.global_step % 1000 == 0 and self.global_rank == 0: |
| 205 | B, N, C, H, W = render_gt['target_images'].shape |
| 206 | N_in = lrm_generator_input['images'].shape[1] |
| 207 | |
| 208 | input_images = v2.functional.resize( |
| 209 | lrm_generator_input['images'], (H, W), interpolation=3, antialias=True).clamp(0, 1) |
| 210 | input_images = torch.cat( |
| 211 | [input_images, torch.ones(B, N-N_in, C, H, W).to(input_images)], dim=1) |
| 212 | |
| 213 | input_images = rearrange( |
| 214 | input_images, 'b n c h w -> b c h (n w)') |
| 215 | target_images = rearrange( |
| 216 | render_gt['target_images'], 'b n c h w -> b c h (n w)') |
| 217 | render_images = rearrange( |
| 218 | render_out['render_images'], 'b n c h w -> b c h (n w)') |
| 219 | target_alphas = rearrange( |
| 220 | repeat(render_gt['target_alphas'], 'b n 1 h w -> b n 3 h w'), 'b n c h w -> b c h (n w)') |
| 221 | render_alphas = rearrange( |
| 222 | repeat(render_out['render_alphas'], 'b n 1 h w -> b n 3 h w'), 'b n c h w -> b c h (n w)') |
| 223 | target_depths = rearrange( |
| 224 | repeat(render_gt['target_depths'], 'b n 1 h w -> b n 3 h w'), 'b n c h w -> b c h (n w)') |
| 225 | render_depths = rearrange( |
| 226 | repeat(render_out['render_depths'], 'b n 1 h w -> b n 3 h w'), 'b n c h w -> b c h (n w)') |
| 227 | MAX_DEPTH = torch.max(target_depths) |
| 228 | target_depths = target_depths / MAX_DEPTH * target_alphas |
| 229 | render_depths = render_depths / MAX_DEPTH |
| 230 | |
| 231 | grid = torch.cat([ |
| 232 | input_images, |
| 233 | target_images, render_images, |
| 234 | target_alphas, render_alphas, |
| 235 | target_depths, render_depths, |
| 236 | ], dim=-2) |
| 237 | grid = make_grid(grid, nrow=target_images.shape[0], normalize=True, value_range=(0, 1)) |
| 238 | |
| 239 | save_image(grid, os.path.join(self.logdir, 'images', f'train_{self.global_step:07d}.png')) |
| 240 | |
| 241 | return loss |
| 242 | |
| 243 | def compute_loss(self, render_out, render_gt): |
| 244 | # NOTE: the rgb value range of OpenLRM is [0, 1] |
nothing calls this directly
no test coverage detected