File size: 2,410 Bytes
23fa981 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
import lightning as L
import torch
from torch import nn
from torchmetrics.functional import accuracy
from torchvision import models
class DRModel(L.LightningModule):
def __init__(
self, num_classes: int, learning_rate: float = 2e-4, class_weights=None
):
super().__init__()
self.save_hyperparameters()
self.num_classes = num_classes
self.learning_rate = learning_rate
# Define the model
# self.model = models.densenet121(weights=models.DenseNet121_Weights.DEFAULT)
self.model = models.densenet169(weights=models.DenseNet169_Weights.DEFAULT)
# self.model = models.vit_b_16(weights=models.ViT_B_16_Weights.DEFAULT)
# freeze the feature extractor
for param in self.model.parameters():
param.requires_grad = False
# Change the output layer to have the number of classes
in_features = self.model.classifier.in_features
# in_features = 768
self.model.classifier = nn.Sequential(
nn.Linear(in_features, in_features // 2),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(in_features // 2, num_classes),
)
# Define the loss function
self.criterion = nn.CrossEntropyLoss(weight=class_weights)
def forward(self, x):
return self.model(x)
def training_step(self, batch):
x, y = batch
logits = self.model(x)
loss = self.criterion(logits, y)
self.log("train_loss", loss, prog_bar=True)
return loss
def validation_step(self, batch, batch_idx):
x, y = batch
logits = self.model(x)
loss = self.criterion(logits, y)
preds = torch.argmax(logits, dim=1)
acc = accuracy(preds, y, task="multiclass", num_classes=self.num_classes)
self.log("val_loss", loss, prog_bar=True)
self.log("val_acc", acc, prog_bar=True)
def configure_optimizers(self):
optimizer = torch.optim.Adam(
self.parameters(), lr=self.learning_rate, weight_decay=1e-4
)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
return {
"optimizer": optimizer,
"lr_scheduler": {
"scheduler": scheduler,
"interval": "epoch",
"monitor": "val_loss",
},
}
# return optimizer
|