Log num params to W&B (#741)

* log num params

* typo

* check logger exists first

* update tb loggers

src/fairchem/core/common/logger.py

@@ -9,6 +9,7 @@
 
 import logging
 from abc import ABC, abstractmethod
+from typing import Any
 
 import torch
 import wandb
@@ -51,6 +52,10 @@ def log_plots(self, plots) -> None:
     def mark_preempting(self) -> None:
         pass
 
+    @abstractmethod
+    def log_summary(self, summary_dict: dict[str, Any]) -> None:
+        pass
+
 
 @registry.register_logger("wandb")
 class WandBLogger(Logger):
@@ -89,6 +94,10 @@ def log_plots(self, plots, caption: str = "") -> None:
         plots = [wandb.Image(x, caption=caption) for x in plots]
         wandb.log({"data": plots})
 
+    def log_summary(self, summary_dict: dict[str, Any]):
+        for k, v in summary_dict.items():
+            wandb.run.summary[k] = v
+
     def mark_preempting(self) -> None:
         wandb.mark_preempting()
 
@@ -114,7 +123,10 @@ def log(self, update_dict, step: int, split: str = ""):
                 self.writer.add_scalar(key, update_dict[key], step)
 
     def mark_preempting(self) -> None:
-        pass
+        logging.warning("mark_preempting for Tensorboard not supported")
 
     def log_plots(self, plots) -> None:
-        pass
+        logging.warning("log_plots for Tensorboard not supported")
+
+    def log_summary(self, summary_dict: dict[str, Any]) -> None:
+        logging.warning("log_summary for Tensorboard not supported")

src/fairchem/core/trainers/base_trainer.py

@@ -430,6 +430,7 @@ def load_model(self) -> None:
 
         if self.logger is not None:
             self.logger.watch(self.model)
+            self.logger.log_summary({"num_params": self.model.num_params})
 
         if distutils.initialized() and not self.config["noddp"]:
             self.model = DistributedDataParallel(self.model, device_ids=[self.device])