From dd3c5fbacf36f79f203a930ad985a90864041b35 Mon Sep 17 00:00:00 2001
From: dcarron <daniel.carron@idiap.ch>
Date: Wed, 5 Jul 2023 14:37:03 +0200
Subject: [PATCH] Update pretrained densenet configs
---
.../configs/models/densenet_pretrained.py | 20 ++++++++++++++-----
1 file changed, 15 insertions(+), 5 deletions(-)
diff --git a/src/ptbench/configs/models/densenet_pretrained.py b/src/ptbench/configs/models/densenet_pretrained.py
index b018a522..f8908fdb 100644
--- a/src/ptbench/configs/models/densenet_pretrained.py
+++ b/src/ptbench/configs/models/densenet_pretrained.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.01}
-
# 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=True
+ criterion,
+ criterion_valid,
+ optimizer,
+ optimizer_configs,
+ pretrained=True,
+ augmentation_transforms=augmentation_transforms,
)
--
GitLab