使用pytorch_lightning进行MNIST数据集的分类

太香了！！！！！！！！！
本文完全参考官方文档

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


transforms = transforms.Compose([transforms.ToTensor(),
                                 transforms.Normalize((0.1317,), (0.3081,))])
mnist_train = MNIST(root=r'./data/mnist', train=True, transform=transforms, download=False)
mnist_train = DataLoader(mnist_train, batch_size=64)


class LitMNIST(LightningModule):
    def __init__(self):
        super().__init__()

        self.layer_1 = nn.Linear(28*28, 128)
        self.layer_2 = nn.Linear(128, 256)
        self.layer_3 = nn.Linear(256, 10)

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

        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)
        return loss
    def configure_optimizers(self):
        return Adam(self.parameters(), lr=0.0001)


model = LitMNIST()
trainer = Trainer()
trainer.fit(model, mnist_train)

a=1