extract data
(self, batch, batch_idx)
| 494 | } |
| 495 | |
| 496 | def training_step(self, batch, batch_idx): |
| 497 | """ extract data """ |
| 498 | data = self.extract_from_batch(batch) |
| 499 | x1_latent = data["x1_latent"] |
| 500 | x0_latent = data["x0_latent"] |
| 501 | |
| 502 | """ cross-attention conditioning """ |
| 503 | if exists(self.cond_stage): |
| 504 | # fetch conditioning from raw batch |
| 505 | conditioning = batch[self.conditioning_key] |
| 506 | conditioning = self.cond_stage(conditioning) |
| 507 | conditioning = self.apply_cond_dropout(conditioning) |
| 508 | else: |
| 509 | conditioning = None |
| 510 | |
| 511 | """ concatenated context """ |
| 512 | if exists(self.context_key): |
| 513 | # fetch context from preprocessed batch |
| 514 | context = data[self.context_key] |
| 515 | else: |
| 516 | context = None |
| 517 | |
| 518 | """ input """ |
| 519 | x_target = x1_latent |
| 520 | |
| 521 | # define x0 |
| 522 | if self.start_from_noise: |
| 523 | x_source = batch.get("noise", torch.randn_like(x1_latent)) |
| 524 | else: |
| 525 | x_source = x0_latent |
| 526 | |
| 527 | # noise x0 |
| 528 | if self.noise_image: |
| 529 | x_source = self.diffusion.q_sample(x_start=x_source, t=self.noising_step) |
| 530 | |
| 531 | """ loss """ |
| 532 | loss = self.forward( |
| 533 | x0=x_source, |
| 534 | x1=x_target, |
| 535 | context=context, |
| 536 | context_ca=conditioning |
| 537 | ) |
| 538 | bs = x_source.shape[0] |
| 539 | self.log("train/loss", loss, on_step=True, on_epoch=True, batch_size=bs) |
| 540 | |
| 541 | """ log statistics """ |
| 542 | if exists(self.ema_model): self.ema_model.update() |
| 543 | if self.stop_training: self.stop_training_method() |
| 544 | if exists(self.lr_scheduler_cfg): self.lr_schedulers().step() |
| 545 | if self.log_grad_norm: |
| 546 | grad_norm = get_grad_norm(self.model) |
| 547 | self.log("train/grad_norm", grad_norm, on_step=True, on_epoch=False, sync_dist=False) |
| 548 | |
| 549 | return loss |
| 550 | |
| 551 | def validation_step(self, batch, batch_idx): |
| 552 | # save samples for visualization |
nothing calls this directly
no test coverage detected