(self, batch, loss_fn, stage)
| 668 | return self.step(batch, l1_loss, "test") |
| 669 | |
| 670 | def step(self, batch, loss_fn, stage): |
| 671 | with torch.set_grad_enabled(stage == "train" or self.hparams.derivative): |
| 672 | pred, deriv = self(batch) |
| 673 | if stage == "test": |
| 674 | self.inference_results['y_pred'].append(pred.squeeze(-1).detach().cpu()) |
| 675 | self.inference_results['y_true'].append(batch.y.squeeze(-1).detach().cpu()) |
| 676 | if self.hparams.derivative: |
| 677 | self.inference_results['dy_pred'].append(deriv.squeeze(-1).detach().cpu()) |
| 678 | self.inference_results['dy_true'].append(batch.dy.squeeze(-1).detach().cpu()) |
| 679 | |
| 680 | loss_y, loss_dy = 0, 0 |
| 681 | if self.hparams.derivative: |
| 682 | if "y" not in batch: |
| 683 | deriv = deriv + pred.sum() * 0 |
| 684 | |
| 685 | loss_dy = loss_fn(deriv, batch.dy) |
| 686 | |
| 687 | if stage in ["train", "val"] and self.hparams.loss_scale_dy < 1: |
| 688 | if self.ema[stage + "_dy"] is None: |
| 689 | self.ema[stage + "_dy"] = loss_dy.detach() |
| 690 | # apply exponential smoothing over batches to dy |
| 691 | loss_dy = ( |
| 692 | self.hparams.loss_scale_dy * loss_dy |
| 693 | + (1 - self.hparams.loss_scale_dy) * self.ema[stage + "_dy"] |
| 694 | ) |
| 695 | self.ema[stage + "_dy"] = loss_dy.detach() |
| 696 | |
| 697 | if self.hparams.force_weight > 0: |
| 698 | self.losses[stage + "_dy"].append(loss_dy.detach()) |
| 699 | |
| 700 | if "y" in batch: |
| 701 | if batch.y.ndim == 1: |
| 702 | batch.y = batch.y.unsqueeze(1) |
| 703 | |
| 704 | loss_y = loss_fn(pred, batch.y) |
| 705 | |
| 706 | if stage in ["train", "val"] and self.hparams.loss_scale_y < 1: |
| 707 | if self.ema[stage + "_y"] is None: |
| 708 | self.ema[stage + "_y"] = loss_y.detach() |
| 709 | # apply exponential smoothing over batches to y |
| 710 | loss_y = ( |
| 711 | self.hparams.loss_scale_y * loss_y |
| 712 | + (1 - self.hparams.loss_scale_y) * self.ema[stage + "_y"] |
| 713 | ) |
| 714 | self.ema[stage + "_y"] = loss_y.detach() |
| 715 | |
| 716 | if self.hparams.energy_weight > 0: |
| 717 | self.losses[stage + "_y"].append(loss_y.detach()) |
| 718 | |
| 719 | loss = loss_y * self.hparams.energy_weight + loss_dy * self.hparams.force_weight |
| 720 | |
| 721 | self.losses[stage].append(loss.detach()) |
| 722 | |
| 723 | return loss |
| 724 | |
| 725 | def optimizer_step(self, *args, **kwargs): |
| 726 | optimizer = kwargs["optimizer"] if "optimizer" in kwargs else args[2] |
no outgoing calls