docker-python/tests/test_pytorch_lightning.py at master · HpcDataLab/docker-python · GitHub

History

68 lines (48 loc) · 1.86 KB

Raw

import unittest

import pytorch_lightning as pl

import torch

import torch.nn.functional as F

from torch.utils.data import DataLoader, TensorDataset

class LitDataModule(pl.LightningDataModule):

def __init__(self, batch_size=16):

super().__init__()

self.batch_size = batch_size

def setup(self, stage=None):

X_train = torch.rand(100, 1, 28, 28)

y_train = torch.randint(0, 10, size=(100,))

X_valid = torch.rand(20, 1, 28, 28)

y_valid = torch.randint(0, 10, size=(20,))

self.train_ds = TensorDataset(X_train, y_train)

self.valid_ds = TensorDataset(X_valid, y_valid)

def train_dataloader(self):

return DataLoader(self.train_ds, batch_size=self.batch_size, shuffle=True)

def val_dataloader(self):

return DataLoader(self.valid_ds, batch_size=self.batch_size, shuffle=False)

class LitClassifier(pl.LightningModule):

def __init__(self):

super().__init__()

self.l1 = torch.nn.Linear(28 * 28, 10)

def forward(self, x):

return F.relu(self.l1(x.view(x.size(0), -1)))

def training_step(self, batch, batch_idx):

x, y = batch

y_hat = self(x)

loss = F.cross_entropy(y_hat, y)

self.log('train_loss', loss)

return loss

def validation_step(self, batch, batch_idx):

x, y = batch

y_hat = self(x)

loss = F.cross_entropy(y_hat, y)

self.log('val_loss', loss)

def configure_optimizers(self):

return torch.optim.Adam(self.parameters(), lr=1e-2)

class TestPytorchLightning(unittest.TestCase):

def test_version(self):

self.assertIsNotNone(pl.__version__)

def test_mnist(self):

dm = LitDataModule()

model = LitClassifier()

trainer = pl.Trainer(gpus=None, max_epochs=1)

result = trainer.fit(model, datamodule=dm)

self.assertTrue(result)

docker-python/tests/test_pytorch_lightning.py at master · HpcDataLab/docker-python · GitHub

Expand file tree

Breadcrumbs

test_pytorch_lightning.py

Latest commit

History

Breadcrumbs

test_pytorch_lightning.py

File metadata and controls