wehub-resource-sync
261 lines
7.2 KiB
Python
261 lines
7.2 KiB
Python
#
|
|
# Trains an MNIST digit recognizer using PyTorch Lightning,
|
|
# and uses MLflow to log metrics, params and artifacts
|
|
# NOTE: This example requires you to first install
|
|
# pytorch-lightning (using pip install pytorch-lightning)
|
|
# and mlflow (using pip install mlflow).
|
|
#
|
|
|
|
|
|
import os
|
|
|
|
import lightning as L
|
|
import torch
|
|
from lightning.pytorch.callbacks import EarlyStopping, LearningRateMonitor, ModelCheckpoint
|
|
from lightning.pytorch.cli import LightningCLI
|
|
from torch.nn import functional as F
|
|
from torch.utils.data import DataLoader, random_split
|
|
from torchvision import datasets, transforms
|
|
|
|
import mlflow.pytorch
|
|
|
|
|
|
class MNISTDataModule(L.LightningDataModule):
|
|
def __init__(self, batch_size=64, num_workers=3):
|
|
"""
|
|
Initialization of inherited lightning data module
|
|
"""
|
|
super().__init__()
|
|
self.df_train = None
|
|
self.df_val = None
|
|
self.df_test = None
|
|
self.train_data_loader = None
|
|
self.val_data_loader = None
|
|
self.test_data_loader = None
|
|
self.batch_size = batch_size
|
|
self.num_workers = num_workers
|
|
|
|
# transforms for images
|
|
self.transform = transforms.Compose([
|
|
transforms.ToTensor(),
|
|
transforms.Normalize((0.1307,), (0.3081,)),
|
|
])
|
|
|
|
def setup(self, stage=None):
|
|
"""
|
|
Downloads the data, parse it and split the data into train, test, validation data
|
|
|
|
Args:
|
|
stage: Stage - training or testing
|
|
"""
|
|
|
|
self.df_train = datasets.MNIST(
|
|
"dataset", download=True, train=True, transform=self.transform
|
|
)
|
|
self.df_train, self.df_val = random_split(self.df_train, [55000, 5000])
|
|
self.df_test = datasets.MNIST(
|
|
"dataset", download=True, train=False, transform=self.transform
|
|
)
|
|
|
|
def create_data_loader(self, df):
|
|
"""
|
|
Generic data loader function
|
|
|
|
Args:
|
|
df: Input tensor
|
|
|
|
Returns:
|
|
Returns the constructed dataloader
|
|
"""
|
|
return DataLoader(df, batch_size=self.batch_size, num_workers=self.num_workers)
|
|
|
|
def train_dataloader(self):
|
|
"""
|
|
Returns:
|
|
output: Train data loader for the given input.
|
|
"""
|
|
return self.create_data_loader(self.df_train)
|
|
|
|
def val_dataloader(self):
|
|
"""
|
|
Returns:
|
|
output: Validation data loader for the given input.
|
|
"""
|
|
return self.create_data_loader(self.df_val)
|
|
|
|
def test_dataloader(self):
|
|
"""
|
|
Returns:
|
|
output: Test data loader for the given input.
|
|
"""
|
|
return self.create_data_loader(self.df_test)
|
|
|
|
|
|
class LightningMNISTClassifier(L.LightningModule):
|
|
def __init__(self, learning_rate=0.01):
|
|
"""
|
|
Initializes the network
|
|
"""
|
|
super().__init__()
|
|
|
|
# mnist images are (1, 28, 28) (channels, width, height)
|
|
self.optimizer = None
|
|
self.scheduler = None
|
|
self.layer_1 = torch.nn.Linear(28 * 28, 128)
|
|
self.layer_2 = torch.nn.Linear(128, 256)
|
|
self.layer_3 = torch.nn.Linear(256, 10)
|
|
self.learning_rate = learning_rate
|
|
self.val_outputs = []
|
|
self.test_outputs = []
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Args:
|
|
x: Input data
|
|
|
|
Returns:
|
|
output - mnist digit label for the input image
|
|
"""
|
|
batch_size = x.size()[0]
|
|
|
|
# (b, 1, 28, 28) -> (b, 1*28*28)
|
|
x = x.view(batch_size, -1)
|
|
|
|
# layer 1 (b, 1*28*28) -> (b, 128)
|
|
x = self.layer_1(x)
|
|
x = torch.relu(x)
|
|
|
|
# layer 2 (b, 128) -> (b, 256)
|
|
x = self.layer_2(x)
|
|
x = torch.relu(x)
|
|
|
|
# layer 3 (b, 256) -> (b, 10)
|
|
x = self.layer_3(x)
|
|
|
|
# probability distribution over labels
|
|
x = torch.log_softmax(x, dim=1)
|
|
|
|
return x
|
|
|
|
def cross_entropy_loss(self, logits, labels):
|
|
"""
|
|
Initializes the loss function
|
|
|
|
Returns:
|
|
output: Initialized cross entropy loss function.
|
|
"""
|
|
return F.nll_loss(logits, labels)
|
|
|
|
def training_step(self, train_batch, batch_idx):
|
|
"""
|
|
Training the data as batches and returns training loss on each batch
|
|
|
|
Args:
|
|
train_batch: Batch data
|
|
batch_idx: Batch indices
|
|
|
|
Returns:
|
|
output - Training loss
|
|
"""
|
|
x, y = train_batch
|
|
logits = self.forward(x)
|
|
loss = self.cross_entropy_loss(logits, y)
|
|
return {"loss": loss}
|
|
|
|
def validation_step(self, val_batch, batch_idx):
|
|
"""
|
|
Performs validation of data in batches
|
|
|
|
Args:
|
|
val_batch: Batch data
|
|
batch_idx: Batch indices
|
|
|
|
Returns:
|
|
output: valid step loss
|
|
"""
|
|
x, y = val_batch
|
|
logits = self.forward(x)
|
|
loss = self.cross_entropy_loss(logits, y)
|
|
self.val_outputs.append(loss)
|
|
return {"val_step_loss": loss}
|
|
|
|
def on_validation_epoch_end(self):
|
|
"""
|
|
Computes average validation loss
|
|
"""
|
|
avg_loss = torch.stack(self.val_outputs).mean()
|
|
self.log("val_loss", avg_loss, sync_dist=True)
|
|
self.val_outputs.clear()
|
|
|
|
def test_step(self, test_batch, batch_idx):
|
|
"""
|
|
Performs test and computes the accuracy of the model
|
|
|
|
Args:
|
|
test_batch: Batch data
|
|
batch_idx: Batch indices
|
|
|
|
Returns:
|
|
output: Testing accuracy
|
|
"""
|
|
x, y = test_batch
|
|
output = self.forward(x)
|
|
_, y_hat = torch.max(output, dim=1)
|
|
test_acc = (y_hat == y).float().mean()
|
|
self.test_outputs.append(test_acc)
|
|
return {"test_acc": test_acc}
|
|
|
|
def on_test_epoch_end(self):
|
|
"""
|
|
Computes average test accuracy score
|
|
"""
|
|
avg_test_acc = torch.stack(self.test_outputs).mean()
|
|
self.log("avg_test_acc", avg_test_acc, sync_dist=True)
|
|
self.test_outputs.clear()
|
|
|
|
def configure_optimizers(self):
|
|
"""
|
|
Initializes the optimizer and learning rate scheduler
|
|
|
|
Returns:
|
|
output: Initialized optimizer and scheduler
|
|
"""
|
|
self.optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
|
|
self.scheduler = {
|
|
"scheduler": torch.optim.lr_scheduler.ReduceLROnPlateau(
|
|
self.optimizer,
|
|
mode="min",
|
|
factor=0.2,
|
|
patience=2,
|
|
min_lr=1e-6,
|
|
),
|
|
"monitor": "val_loss",
|
|
}
|
|
return [self.optimizer], [self.scheduler]
|
|
|
|
|
|
def cli_main():
|
|
early_stopping = EarlyStopping(
|
|
monitor="val_loss",
|
|
)
|
|
|
|
checkpoint_callback = ModelCheckpoint(
|
|
dirpath=os.getcwd(), save_top_k=1, verbose=True, monitor="val_loss", mode="min"
|
|
)
|
|
lr_logger = LearningRateMonitor()
|
|
cli = LightningCLI(
|
|
LightningMNISTClassifier,
|
|
MNISTDataModule,
|
|
run=False,
|
|
save_config_callback=None,
|
|
trainer_defaults={"callbacks": [early_stopping, checkpoint_callback, lr_logger]},
|
|
)
|
|
if cli.trainer.global_rank == 0:
|
|
mlflow.pytorch.autolog()
|
|
cli.trainer.fit(cli.model, datamodule=cli.datamodule)
|
|
cli.trainer.test(ckpt_path="best", datamodule=cli.datamodule)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cli_main()
|