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/configs/models/densenet.py

@@ -6,20 +6,30 @@
 
 from torch import empty
 from torch.nn import BCEWithLogitsLoss
+from torch.optim import Adam
 
 from ...models.densenet import Densenet
 
-# config
-optimizer_configs = {"lr": 0.0001}
-
 # optimizer
-optimizer = "Adam"
+optimizer = Adam
+optimizer_configs = {"lr": 0.0001}
 
 # criterion
 criterion = BCEWithLogitsLoss(pos_weight=empty(1))
 criterion_valid = BCEWithLogitsLoss(pos_weight=empty(1))
 
+from ...data.transforms import ElasticDeformation
+
+augmentation_transforms = [
+    ElasticDeformation(p=0.8),
+]
+
 # model
 model = Densenet(
-    criterion, criterion_valid, optimizer, optimizer_configs, pretrained=False
+    criterion,
+    criterion_valid,
+    optimizer,
+    optimizer_configs,
+    pretrained=False,
+    augmentation_transforms=augmentation_transforms,
 )