additional logging

JVGD · Nov 1, 2021 · 456b02d · 456b02d
1 parent eebaf4e
commit 456b02d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/experiment1/model.py b/experiment1/model.py
@@ -56,4 +56,8 @@ def validation_step(self, batch, batch_idx, *args, **kwargs) -> None:
         self.log("loss/valid", loss)
 
     def configure_optimizers(self):
-        return optim.SGD(self.parameters(), lr=self.lr, **self.optim_conf)
+        return optim.SGD(self.parameters(), lr=self.lr, **self.optim_conf)
+
+    def on_epoch_end(self) -> None:
+        self.log("weights/weights_adder", self.weights_adder)
+        self.log("weights/weights_substractor", self.weights_substracter)
diff --git a/experiment1/trainer.py b/experiment1/trainer.py
@@ -12,6 +12,6 @@
 model = Model1(lr=1e-4, optim_conf={"momentum": 0.9})
 
 # Trainer
-trainer = Trainer(max_epochs=500, progress_bar_refresh_rate=20)
+trainer = Trainer(max_epochs=500)
 trainer.fit(model=model, train_dataloaders=dl_train, val_dataloaders=dl_valid)