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

Moved alexnet_pretrained to lightning

parent 709304f8
No related branches found
No related tags found
1 merge request!4Moved code to lightning
...@@ -2,24 +2,23 @@ ...@@ -2,24 +2,23 @@
# #
# SPDX-License-Identifier: GPL-3.0-or-later # SPDX-License-Identifier: GPL-3.0-or-later
"""AlexNet. """AlexNet."""
Pretrained AlexNet
"""
from torch import empty
from torch.nn import BCEWithLogitsLoss from torch.nn import BCEWithLogitsLoss
from torch.optim import SGD
from ...models.alexnet import build_alexnet from ...models.alexnet import Alexnet
# config # config
lr = 0.001 optimizer_configs = {"lr": 0.001, "momentum": 0.1}
# model
model = build_alexnet(pretrained=True)
# optimizer # optimizer
optimizer = SGD(model.parameters(), lr=lr, momentum=0.1) optimizer = "SGD"
# criterion # criterion
criterion = BCEWithLogitsLoss() criterion = BCEWithLogitsLoss(pos_weight=empty(1))
criterion_valid = BCEWithLogitsLoss(pos_weight=empty(1))
# model
model = Alexnet(
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