From c5f71ad472964e35ae4c83d54dfc3c3b79b84c6c Mon Sep 17 00:00:00 2001
From: dcarron <daniel.carron@idiap.ch>
Date: Tue, 11 Apr 2023 13:02:16 +0200
Subject: [PATCH] Moved densenet_pretrained to lightning
---
.../configs/models/densenet_pretrained.py | 19 +++++++++++--------
1 file changed, 11 insertions(+), 8 deletions(-)
diff --git a/src/ptbench/configs/models/densenet_pretrained.py b/src/ptbench/configs/models/densenet_pretrained.py
index cb729083..b018a522 100644
--- a/src/ptbench/configs/models/densenet_pretrained.py
+++ b/src/ptbench/configs/models/densenet_pretrained.py
@@ -4,19 +4,22 @@
"""DenseNet."""
+from torch import empty
from torch.nn import BCEWithLogitsLoss
-from torch.optim import Adam
-from ...models.densenet import build_densenet
+from ...models.densenet import Densenet
# config
-lr = 0.01
-
-# model
-model = build_densenet(pretrained=True)
+optimizer_configs = {"lr": 0.01}
# optimizer
-optimizer = Adam(model.parameters(), lr=lr)
+optimizer = "Adam"
# criterion
-criterion = BCEWithLogitsLoss()
+criterion = BCEWithLogitsLoss(pos_weight=empty(1))
+criterion_valid = BCEWithLogitsLoss(pos_weight=empty(1))
+
+# model
+model = Densenet(
+ criterion, criterion_valid, optimizer, optimizer_configs, pretrained=True
+)
--
GitLab