pytorch_lightning

1. train

import torch
from torch.nn import functional as F
from torch import nn
from pytorch_lightning.core.lightning import LightningModule
import pytorch_lightning as pl

import os
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, random_split
from torchvision.datasets import MNIST
from torchvision import datasets, transforms
import pytorch_lightning as pl
from pytorch_lightning import Trainer
from pytorch_lightning.core.lightning import LightningModule


class MNISTDataModule(pl.LightningDataModule):
    def __init__(self, batch_size=8):
        super().__init__()
        self.batch_size = batch_size

    def setup(self, stage):
        # transform
        transform = transforms.Compose([transforms.ToTensor()])
        mnist_train = MNIST('./1/mnist/', train=True, download=True, transform=transform)
        mnist_test = MNIST('./1/mnist/', train=False, download=True, transform=transform)

        # train/val split
        mnist_train, mnist_val = random_split(mnist_train, [55000, 5000])

        # assign to use in dataloaders
        self.train_dataset = mnist_train
        self.val_dataset = mnist_val
        self.test_dataset = mnist_test

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size)

    def val_dataloader(self):
        return DataLoader(self.val_dataset, batch_size=self.batch_size)

    def test_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=self.batch_size)


class LitMNIST(LightningModule):

    def __init__(self):
        super().__init__()

        # mnist images are (1, 28, 28) (channels, width, height)
        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)

    def forward(self, x):
        batch_size, channels, width, height = x.size()

        # (b, 1, 28, 28) -> (b, 1*28*28)
        x = x.view(batch_size, -1)
        x = self.layer_1(x)
        x = F.relu(x)
        x = self.layer_2(x)
        x = F.relu(x)
        x = self.layer_3(x)

        x = F.log_softmax(x, dim=1)
        return x

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        print(self.global_step)
        self.logger.experiment.add_images('train/imgs', x, self.global_step)
        self.log('train/global_step', self.global_step, prog_bar=False)
        self.log('train/loss', loss, prog_bar=False)
        return loss

    def validation_step(self, batch, batch_nb):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        print(self.global_step)
        self.logger.experiment.add_images(f'val/imgs/{self.global_step}', x, batch_nb)
        self.log('val/global_step', self.global_step, prog_bar=False)
        self.log('val/loss', loss, prog_bar=False)
        return [loss, x]

    def validation_epoch_end(self, outputs):
        imgs = []
        for i, [loss, x] in enumerate(outputs):
            imgs.append(x)
            print(x.shape)

        imgs = torch.stack(imgs, 0)
        imgs = imgs.reshape(-1, imgs.shape[2], imgs.shape[3], imgs.shape[4])
        self.logger.experiment.add_images(f'val/imgs', imgs, self.global_step)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

    def infer(self,x):
        self.eval()
        x = self.forward(x)
        return x


if __name__ == "__main__":
    from pytorch_lightning import Trainer
    dm = MNISTDataModule()
    model = LitMNIST()
    trainer = Trainer(gpus=1,max_epochs=20)
    trainer.fit(model, dm)

2. eval

from pytorch_lightning.core.lightning import LightningModule
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, random_split
from torchvision.datasets import MNIST
from torchvision import  transforms
from minist_base_1 import LitMNIST

transform = transforms.Compose([transforms.ToTensor()])
mnist_train = MNIST('./1/mnist/', train=True, download=True, transform=transform)
mnist_test = MNIST('./1/mnist/', train=False, download=True, transform=transform)

mnist_train, mnist_val = random_split(mnist_train, [55000, 5000])

train_dataset = mnist_train
val_dataset = mnist_val
test_dataset = mnist_test

def train_dataloader():
    return DataLoader(train_dataset, batch_size=32)

def val_dataloader():
    return DataLoader(val_dataset, batch_size=2)

def test_dataloader():
    return DataLoader(test_dataset, batch_size=1)

if __name__ == "__main__":
    model = LitMNIST()
    model = model.load_from_checkpoint("lightning_logs/version_2/checkpoints/epoch=0-step=6874.ckpt")
    #model = model.load_from_checkpoint("lightning_logs/version_1/checkpoints/epoch=5-step=41249.ckpt")
    count = 0
    for ii,data in enumerate(test_dataset):
        x,y = data
        x = x.unsqueeze(0)
        predict = model(x)
        if torch.argmax(predict).cpu().data!=y:
            count += 1
            print(count,torch.argmax(predict).item(),y)