Skip to content
Snippets Groups Projects
Commit c5f71ad4 authored by Daniel CARRON's avatar Daniel CARRON :b:
Browse files

Moved densenet_pretrained to lightning

parent 1099d715
No related branches found
No related tags found
1 merge request!4Moved code to lightning
......@@ -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
)
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment