Skip to content
Snippets Groups Projects

Making use of LightningDataModule and simplification of data loading

Merged Daniel CARRON requested to merge add-datamodule into main
Compare and Show latest version
15 files
+ 224
220
Compare changes
  • Side-by-side
  • Inline
Files
15
@@ -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,
)
Loading