add support for different inputs - either raw batches of the output o…

…f PVNet
openclimatefix · Jul 20, 2023 · 6011359 · 6011359
1 parent 3b801ee
commit 6011359
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/pvnet_summation/models/base_model.py b/pvnet_summation/models/base_model.py
@@ -135,7 +135,8 @@ def _training_accumulate_log(self, batch_idx, losses, y_hat, y, times):
 
     def training_step(self, batch, batch_idx):
         """Run training step"""
-        y_hat = self.forward(batch['pvnet_inputs'])
+
+        y_hat = self.forward(batch)
         y = batch["national_targets"]
         times = batch["times"]
 
@@ -152,7 +153,8 @@ def training_step(self, batch, batch_idx):
 
     def validation_step(self, batch: dict, batch_idx):
         """Run validation step"""
-        y_hat = self.forward(batch['pvnet_inputs'])
+
+        y_hat = self.forward(batch)
         y = batch["national_targets"]
         times = batch["times"]
 
@@ -201,7 +203,8 @@ def validation_step(self, batch: dict, batch_idx):
 
     def test_step(self, batch, batch_idx):
         """Run test step"""
-        y_hat = self.forward(batch['pvnet_inputs'])
+
+        y_hat = self.forward(batch)
         y = batch["national_targets"]
 
         losses = self._calculate_common_losses(y, y_hat)

diff --git a/pvnet_summation/models/model.py b/pvnet_summation/models/model.py
@@ -66,8 +66,13 @@ def __init__(
 
 
     def forward(self, x):
-        """Run central model forward"""
-        pvnet_out = self.predict_pvnet_batch(x)
+        """Run model forward"""
+
+        if "pvnet_outputs" in x:
+            pvnet_out = x["pvnet_outputs"]
+        else:
+            pvnet_out = self.predict_pvnet_batch(x['pvnet_inputs'])
+
         pvnet_out = torch.flatten(pvnet_out, start_dim=1)
         out = self.model(pvnet_out)