From 7df06c45b7f5330f59bd9ced58d426050e75650d Mon Sep 17 00:00:00 2001
From: Theo Heimel <info@theovention.de>
Date: Thu, 16 Mar 2023 14:21:32 +0100
Subject: [PATCH] various changes

---
 README.md    |  3 ++-
 src/model.py |  2 +-
 src/plots.py |  6 +++---
 src/train.py | 10 ++++++++--
 4 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 23dbb4b..a7a76b4 100644
--- a/README.md
+++ b/README.md
@@ -47,7 +47,7 @@ Parameter        | Description
 ### Training
 
 Parameter         | Description
-------------------|--------------------------------------------------------
+------------------|---------------------------------------------------------------------------
 `bayesian`        | Train as a Bayesian network
 `batch_size`      | Batch size
 `lr`              | Initial learning rate
@@ -59,6 +59,7 @@ Parameter         | Description
 `lr_decay_epochs` | Step scheduler: Epochs after which to reduce the LR
 `lr_decay_factor` | Step scheduler: Decay factor
 `epochs`          | Number of epochs
+`train_samples`   | Total number of samples used for training (alternative to number of epochs)
 
 ### Evaluation
 
diff --git a/src/model.py b/src/model.py
index bdcf1ee..c542cb9 100644
--- a/src/model.py
+++ b/src/model.py
@@ -45,7 +45,7 @@ def __init__(self, input_dim: int, params: dict):
                 self.bayesian_layers.append(layer)
             layers.append(layer)
             if dropout > 0:
-                self.layers.append(nn.Dropout(p=dropout))
+                layers.append(nn.Dropout(p=dropout))
             layers.append(activation())
             layer_size = hidden_size
         layer = layer_class(layer_size, 1, **layer_kwargs)
diff --git a/src/plots.py b/src/plots.py
index 06b588a..bd46aa2 100644
--- a/src/plots.py
+++ b/src/plots.py
@@ -276,7 +276,7 @@ def plot_single_weight_hist(
             ax,
             bins,
             y_combined / np.sum(y_combined),
-            y_combined_err / np.sum(y_combined),
+            y_combined_err / np.sum(y_combined) if y_combined_err is not None else None,
             label = "Comb",
             color = self.colors[0]
         )
@@ -284,7 +284,7 @@ def plot_single_weight_hist(
             ax,
             bins,
             y_true / np.sum(y_true),
-            y_true_err / np.sum(y_true),
+            y_true_err / np.sum(y_true) if y_true_err is not None else None,
             label = "Truth",
             color = self.colors[1]
         )
@@ -292,7 +292,7 @@ def plot_single_weight_hist(
             ax,
             bins,
             y_fake / np.sum(y_fake),
-            y_fake_err / np.sum(y_fake),
+            y_fake_err / np.sum(y_fake) if y_fake_err is not None else None,
             label = "Gen",
             color = self.colors[2]
         )
diff --git a/src/train.py b/src/train.py
index 9d7648e..d9af34b 100644
--- a/src/train.py
+++ b/src/train.py
@@ -94,6 +94,12 @@ def init_scheduler(self):
         """
         Initialized the LR scheduler. Currently, one-cycle and step schedulers are supported.
         """
+        if "epochs" in self.params:
+            self.epochs = self.params["epochs"]
+        else:
+            self.epochs = int(
+                self.params["train_samples"] / self.params["batch_size"] / self.train_batches
+            )
         self.scheduler_type = self.params.get("lr_scheduler", "one_cycle")
         if self.scheduler_type == "step":
             self.scheduler = torch.optim.lr_scheduler.StepLR(
@@ -105,7 +111,7 @@ def init_scheduler(self):
             self.scheduler = torch.optim.lr_scheduler.OneCycleLR(
                 self.optimizer,
                 self.params.get("max_lr", self.params["lr"]*10),
-                epochs = self.params["epochs"],
+                epochs = self.epochs,
                 steps_per_epoch=self.train_batches
             )
         else:
@@ -144,7 +150,7 @@ def train(self):
         """
         Main training loop
         """
-        for epoch in range(self.params["epochs"]):
+        for epoch in range(self.epochs):
             self.model.train()
             epoch_losses, epoch_bce_losses, epoch_kl_losses = [], [], []
             for (x_true, ), (x_fake, ) in zip(self.train_loader_true, self.train_loader_fake):