wehub-resource-sync
172 lines
5.6 KiB
Python
172 lines
5.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torch.utils.data import DataLoader
|
|
from torchmetrics import Accuracy
|
|
|
|
from ray import train
|
|
from ray.train.lightning._lightning_utils import import_lightning
|
|
|
|
pl = import_lightning()
|
|
|
|
|
|
class LinearModule(pl.LightningModule):
|
|
def __init__(self, input_dim, output_dim, strategy="ddp", fail_epoch=-1) -> None:
|
|
super().__init__()
|
|
self.linear = nn.Linear(input_dim, output_dim)
|
|
self.loss = []
|
|
self.strategy = strategy
|
|
self.restored = train.get_checkpoint() is not None
|
|
self.fail_epoch = fail_epoch
|
|
|
|
def forward(self, input):
|
|
if isinstance(input, dict) and len(input) == 1:
|
|
input = list(input.values())[0]
|
|
return self.linear(input)
|
|
|
|
def training_step(self, batch):
|
|
if not self.restored and self.fail_epoch == self.current_epoch:
|
|
raise RuntimeError
|
|
|
|
output = self.forward(batch)
|
|
loss = torch.sum(output)
|
|
self.log("loss", loss)
|
|
return loss
|
|
|
|
def validation_step(self, val_batch, batch_idx):
|
|
loss = self.forward(val_batch)
|
|
self.loss.append(loss)
|
|
return {"val_loss": loss}
|
|
|
|
def test_step(self, batch, batch_idx):
|
|
loss = self.forward(batch)
|
|
return {"test_loss": loss}
|
|
|
|
def on_validation_epoch_end(self) -> None:
|
|
avg_loss = torch.stack(self.loss).mean()
|
|
self.log("val_loss", avg_loss)
|
|
self.loss.clear()
|
|
|
|
def predict_step(self, batch, batch_idx):
|
|
return self.forward(batch)
|
|
|
|
def configure_optimizers(self):
|
|
if self.strategy == "fsdp":
|
|
# Feed FSDP wrapped model parameters to optimizer
|
|
return torch.optim.AdamW(self.trainer.model.parameters(), lr=0.1)
|
|
else:
|
|
return torch.optim.AdamW(self.parameters(), lr=0.1)
|
|
|
|
|
|
class DoubleLinearModule(pl.LightningModule):
|
|
def __init__(self, input_dim_1, input_dim_2, output_dim) -> None:
|
|
super().__init__()
|
|
self.linear_1 = nn.Linear(input_dim_1, output_dim)
|
|
self.linear_2 = nn.Linear(input_dim_2, output_dim)
|
|
self.loss = []
|
|
|
|
def forward(self, batch):
|
|
input_1 = batch["input_1"]
|
|
input_2 = batch["input_2"]
|
|
return self.linear_1(input_1) + self.linear_2(input_2)
|
|
|
|
def training_step(self, batch):
|
|
output = self.forward(batch)
|
|
loss = torch.sum(output)
|
|
self.log("loss", loss)
|
|
return loss
|
|
|
|
def validation_step(self, val_batch, batch_idx):
|
|
loss = self.forward(val_batch)
|
|
self.loss.append(loss)
|
|
return {"val_loss": loss}
|
|
|
|
def on_validation_epoch_end(self) -> None:
|
|
print("Validation Epoch:", self.current_epoch)
|
|
avg_loss = torch.stack(self.loss).mean()
|
|
self.log("val_loss", avg_loss)
|
|
self.loss.clear()
|
|
|
|
def predict_step(self, batch, batch_idx):
|
|
return self.forward(batch)
|
|
|
|
def configure_optimizers(self):
|
|
return torch.optim.AdamW(self.parameters(), lr=0.1)
|
|
|
|
|
|
class DummyDataModule(pl.LightningDataModule):
|
|
def __init__(self, batch_size: int = 8, dataset_size: int = 256) -> None:
|
|
super().__init__()
|
|
self.batch_size = batch_size
|
|
self.train_data = torch.randn(dataset_size, 32)
|
|
self.val_data = torch.randn(dataset_size, 32)
|
|
self.test_data = torch.randn(dataset_size, 32)
|
|
|
|
def train_dataloader(self):
|
|
return DataLoader(self.train_data, batch_size=self.batch_size)
|
|
|
|
def val_dataloader(self):
|
|
return DataLoader(self.val_data, batch_size=self.batch_size)
|
|
|
|
def test_dataloader(self):
|
|
return DataLoader(self.test_data, batch_size=self.batch_size)
|
|
|
|
|
|
class LightningMNISTClassifier(pl.LightningModule):
|
|
def __init__(self, lr: float, layer_1: int, layer_2: int):
|
|
super(LightningMNISTClassifier, self).__init__()
|
|
|
|
self.lr = lr
|
|
# mnist images are (1, 28, 28) (channels, width, height)
|
|
self.layer_1 = torch.nn.Linear(28 * 28, layer_1)
|
|
self.layer_2 = torch.nn.Linear(layer_1, layer_2)
|
|
self.layer_3 = torch.nn.Linear(layer_2, 10)
|
|
self.accuracy = Accuracy(task="multiclass", num_classes=10, top_k=1)
|
|
self.val_acc_list = []
|
|
self.val_loss_list = []
|
|
|
|
def forward(self, x):
|
|
batch_size, channels, width, height = x.size()
|
|
x = x.view(batch_size, -1)
|
|
x = self.layer_1(x)
|
|
x = torch.relu(x)
|
|
x = self.layer_2(x)
|
|
x = torch.relu(x)
|
|
x = self.layer_3(x)
|
|
x = torch.log_softmax(x, dim=1)
|
|
return x
|
|
|
|
def configure_optimizers(self):
|
|
return torch.optim.Adam(self.parameters(), lr=self.lr)
|
|
|
|
def training_step(self, train_batch, batch_idx):
|
|
x, y = train_batch
|
|
logits = self.forward(x)
|
|
loss = F.nll_loss(logits, y)
|
|
acc = self.accuracy(logits, y)
|
|
self.log("ptl/train_loss", loss)
|
|
self.log("ptl/train_accuracy", acc)
|
|
return loss
|
|
|
|
def validation_step(self, val_batch, batch_idx):
|
|
x, y = val_batch
|
|
logits = self.forward(x)
|
|
loss = F.nll_loss(logits, y)
|
|
acc = self.accuracy(logits, y)
|
|
self.val_acc_list.append(acc)
|
|
self.val_loss_list.append(loss)
|
|
return {"val_loss": loss, "val_accuracy": acc}
|
|
|
|
def on_validation_epoch_end(self):
|
|
avg_loss = torch.stack(self.val_loss_list).mean()
|
|
avg_acc = torch.stack(self.val_acc_list).mean()
|
|
self.log("ptl/val_loss", avg_loss)
|
|
self.log("ptl/val_accuracy", avg_acc)
|
|
self.val_acc_list.clear()
|
|
self.val_loss_list.clear()
|
|
|
|
def predict_step(self, batch, batch_idx, dataloader_idx=None):
|
|
x = batch
|
|
logits = self.forward(x)
|
|
return torch.argmax(logits, dim=-1)
|