File size: 3,004 Bytes
c3d82b0 |
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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import lightning as pl
import torch.nn as nn
import torch
from timm import create_model
from torchmetrics.classification import Accuracy
from torch.optim.lr_scheduler import StepLR
import torch.optim as optim
from loguru import logger
logger.add("logs/model.log", rotation="1 MB", level="INFO")
class LitEfficientNet(pl.LightningModule):
def __init__(
self,
model_name="tf_efficientnet_lite0",
num_classes=10,
lr=1e-3,
custom_loss=None,
):
"""
Initializes a CNN model from TIMM and integrates TorchMetrics.
Args:
model_name (str): TIMM model name (e.g., "tf_efficientnet_lite0").
num_classes (int): Number of output classes (e.g., 0–9 for MNIST).
lr (float): Learning rate for the optimizer.
custom_loss (callable, optional): Custom loss function. Defaults to CrossEntropyLoss.
"""
super().__init__()
self.lr = lr
self.model = create_model(
model_name,
pretrained=True,
num_classes=num_classes,
in_chans=1, # Set to 1 channel for grayscale input
)
self.loss_fn = custom_loss or nn.CrossEntropyLoss()
self.train_acc = Accuracy(num_classes=num_classes, task="multiclass")
self.val_acc = Accuracy(num_classes=num_classes, task="multiclass")
self.test_acc = Accuracy(num_classes=num_classes, task="multiclass")
logger.info(f"Model initialized with TIMM backbone: {model_name}")
logger.info(f"Number of output classes: {num_classes}")
def forward(self, x):
"""
Forward pass of the model.
Args:
x (torch.Tensor): Input tensor.
Returns:
torch.Tensor: Model predictions.
"""
return self.model(x)
def training_step(self, batch, batch_idx):
x, y = batch
y_hat = self(x)
loss = self.loss_fn(y_hat, y)
self.train_acc.update(y_hat, y)
self.log("train_loss", loss, prog_bar=True, logger=True)
self.log("train_acc", self.train_acc, prog_bar=True, logger=True)
return loss
def validation_step(self, batch, batch_idx):
x, y = batch
y_hat = self(x)
loss = self.loss_fn(y_hat, y)
self.val_acc.update(y_hat, y)
self.log("val_loss", loss, prog_bar=True, logger=True)
self.log("val_acc", self.val_acc, prog_bar=True, logger=True)
def test_step(self, batch, batch_idx):
x, y = batch
y_hat = self(x)
self.test_acc.update(y_hat, y)
self.log("test_acc", self.test_acc, prog_bar=True, logger=True)
def configure_optimizers(self):
optimizer = optim.Adam(self.parameters(), lr=self.lr)
scheduler = StepLR(optimizer, step_size=1, gamma=0.9)
logger.info(f"Optimizer: Adam, Learning Rate: {self.lr}")
logger.info("Scheduler: StepLR with step_size=1 and gamma=0.9")
return [optimizer], [scheduler]
|