Reviewed DataModule design+docs+types

In this MR, I try to review the use of the data module and go further in simplifications for the whole code:

• Slightly streamlined the datamodule approach
• Added documentation
• Added type annotations
• Added some TODOs for further discussion

Closes #7 (closed).

TODO:

• (@dcarron) Implement/document/type hint normaliser (ensure it is used correctly through all the models; integrate on the train script) and that train CLI works!
• (@andre.anjos) Deal with import reweight_BCEWithLogitsLoss

src/ptbench/engine/callbacks.py

@@ -95,7 +95,7 @@ class LoggingCallback(Callback):
             self.log("eta", eta_seconds)
             self.log("loss", numpy.average(self.training_loss))
             self.log("learning_rate", pl_module.optimizer_configs["lr"])
-            self.log("validation_loss", numpy.sum(self.validation_loss))
+            self.log("validation_loss", numpy.average(self.validation_loss))
 
             if len(self.extra_validation_loss) > 0:
                 for (
@@ -103,7 +103,8 @@ class LoggingCallback(Callback):
                     extra_valid_loss_values,
                 ) in self.extra_validation_loss.items:
                     self.log(
-                        extra_valid_loss_key, numpy.sum(extra_valid_loss_values)
+                        extra_valid_loss_key,
+                        numpy.average(extra_valid_loss_values),
                     )
 
         queue_retries = 0