make valid loss calculation way same as training

Author: OuyangWenyu

torchhydro/datasets/data_sets.py

@@ -300,7 +300,7 @@ def times(self):
         return times_
 
     def __len__(self):
-        return self.num_samples if self.train_mode else self.ngrid
+        return self.num_samples
 
     def __getitem__(self, item: int):
         basin, idx = self.lookup_table[item]
@@ -818,9 +818,6 @@ def __getitem__(self, index):
         y = ys[-1, :]
         return xc, y
 
-    def __len__(self):
-        return self.num_samples
-
 
 class DplDataset(BaseDataset):
     """pytorch dataset for Differential parameter learning"""
@@ -886,9 +883,6 @@ def __getitem__(self, item):
             z_train,
         ), torch.from_numpy(y_train).float()
 
-    def __len__(self):
-        return self.num_samples if self.train_mode else len(self.t_s_dict["sites_id"])
-
 
 class FlexibleDataset(BaseDataset):
     """A dataset whose datasources are from multiple sources according to the configuration"""
@@ -999,9 +993,6 @@ def _normalize(self):
         # TODO: this work for minio? maybe better to move to basedataset
         return x.compute(), y.compute(), c.compute()
 
-    def __len__(self):
-        return self.num_samples
-
     def __getitem__(self, item: int):
         basin, time = self.lookup_table[item]
         rho = self.rho

torchhydro/trainers/train_utils.py

@@ -482,18 +482,20 @@ def compute_validation(
     seq_first = kwargs["which_first_tensor"] != "batch"
     obs = []
     preds = []
+    valid_loss = 0.0
     with torch.no_grad():
         for src, trg in data_loader:
             trg, output = model_infer(seq_first, device, model, src, trg)
             obs.append(trg)
             preds.append(output)
+            valid_loss_ = compute_loss(trg, output, criterion)
+            valid_loss = valid_loss + valid_loss_.item()
             # clear memory to save GPU memory
             torch.cuda.empty_cache()
         # first dim is batch
         obs_final = torch.cat(obs, dim=0)
         pred_final = torch.cat(preds, dim=0)
-
-        valid_loss = compute_loss(obs_final, pred_final, criterion)
+    valid_loss = valid_loss / len(data_loader)
     y_obs = obs_final.detach().cpu().numpy()
     y_pred = pred_final.detach().cpu().numpy()
     return y_obs, y_pred, valid_loss