Update model configs

src/ptbench/configs/models/alexnet.py

@@ -4,32 +4,17 @@
 
 """AlexNet."""
 
-from torch import empty
 from torch.nn import BCEWithLogitsLoss
 from torch.optim import SGD
 
-from ...models.alexnet import Alexnet
-
-# optimizer
-optimizer = SGD
-optimizer_configs = {"lr": 0.01, "momentum": 0.1}
-
-# criterion
-criterion = BCEWithLogitsLoss(pos_weight=empty(1))
-criterion_valid = BCEWithLogitsLoss(pos_weight=empty(1))
-
 from ...data.transforms import ElasticDeformation
+from ...models.alexnet import Alexnet
 
-augmentation_transforms = [
-    ElasticDeformation(p=0.8),
-]
-
-# model
 model = Alexnet(
-    criterion,
-    criterion_valid,
-    optimizer,
-    optimizer_configs,
+    train_loss=BCEWithLogitsLoss(),
+    validation_loss=BCEWithLogitsLoss(),
+    optimizer_type=SGD,
+    optimizer_arguments=dict(lr=0.01, momentum=0.1),
+    augmentation_transforms=[ElasticDeformation(p=0.8)],
     pretrained=False,
-    augmentation_transforms=augmentation_transforms,
 )

src/ptbench/configs/models/alexnet_pretrained.py

@@ -4,32 +4,17 @@
 
 """AlexNet."""
 
-from torch import empty
 from torch.nn import BCEWithLogitsLoss
 from torch.optim import SGD
 
-from ...models.alexnet import Alexnet
-
-# optimizer
-optimizer = SGD
-optimizer_configs = {"lr": 0.01, "momentum": 0.1}
-
-# criterion
-criterion = BCEWithLogitsLoss(pos_weight=empty(1))
-criterion_valid = BCEWithLogitsLoss(pos_weight=empty(1))
-
 from ...data.transforms import ElasticDeformation
+from ...models.alexnet import Alexnet
 
-augmentation_transforms = [
-    ElasticDeformation(p=0.8),
-]
-
-# model
 model = Alexnet(
-    criterion,
-    criterion_valid,
-    optimizer,
-    optimizer_configs,
+    train_loss=BCEWithLogitsLoss(),
+    validation_loss=BCEWithLogitsLoss(),
+    optimizer_type=SGD,
+    optimizer_arguments=dict(lr=0.01, momentum=0.1),
+    augmentation_transforms=[ElasticDeformation(p=0.8)],
     pretrained=True,
-    augmentation_transforms=augmentation_transforms,
 )

src/ptbench/configs/models/densenet.py

@@ -4,32 +4,17 @@
 
 """DenseNet."""
 
-from torch import empty
 from torch.nn import BCEWithLogitsLoss
 from torch.optim import Adam
 
-from ...models.densenet import Densenet
-
-# optimizer
-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
+from ...models.densenet import Densenet
 
-augmentation_transforms = [
-    ElasticDeformation(p=0.8),
-]
-
-# model
 model = Densenet(
-    criterion,
-    criterion_valid,
-    optimizer,
-    optimizer_configs,
+    train_loss=BCEWithLogitsLoss(),
+    validation_loss=BCEWithLogitsLoss(),
+    optimizer_type=Adam,
+    optimizer_arguments=dict(lr=0.0001),
+    augmentation_transforms=[ElasticDeformation(p=0.8)],
     pretrained=False,
-    augmentation_transforms=augmentation_transforms,
 )

src/ptbench/configs/models/densenet_pretrained.py

@@ -4,32 +4,17 @@
 
 """DenseNet."""
 
-from torch import empty
 from torch.nn import BCEWithLogitsLoss
 from torch.optim import Adam
 
-from ...models.densenet import Densenet
-
-# optimizer
-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
+from ...models.densenet import Densenet
 
-augmentation_transforms = [
-    ElasticDeformation(p=0.8),
-]
-
-# model
 model = Densenet(
-    criterion,
-    criterion_valid,
-    optimizer,
-    optimizer_configs,
+    train_loss=BCEWithLogitsLoss(),
+    validation_loss=BCEWithLogitsLoss(),
+    optimizer_type=Adam,
+    optimizer_arguments=dict(lr=0.0001),
+    augmentation_transforms=[ElasticDeformation(p=0.8)],
     pretrained=True,
-    augmentation_transforms=augmentation_transforms,
 )