Base

所有的代码以官方的MNIST进行测试，扩展

  1. from torch.nn import functional as F
  2. from torch import nn
  3. from torch.utils.data import DataLoader
  4. from torchvision.datasets import MNIST
  5. import os
  6. from torchvision import transforms
  7. from torch.optim import Adam
  8. from pytorch_lightning import Trainer
  9. from pytorch_lightning import LightningModule
  11. class LitMNIST(LightningModule):
  12.     def __init__(self, hparams):
  13.         super().__init__()
  14.         self.layer_1 = nn.Linear(28 * 28, 10)
  15.         self.params = hparams
  16.     def forward(self, x):
  17.         batch_size, channels, height, width = x.size()
  18.         x = x.view(batch_size, -1)
  19.         x = self.layer_1(x)
  20.         x = F.log_softmax(x, dim=1)
  21.         return x
  23.     def training_step(self, batch, batch_idx):  
  24.         x, y = batch
  25.         logits = self(x)
  26.         loss = F.nll_loss(logits, y)
  27.         return loss
  29.     def configure_optimizers(self):
  30.         optimizer = Adam(self.parameters(), lr=self.params['lr'], weight_decay= self.params['weight_decay'])
  31.         return optimizer
  33. hparams={
  34.     'lr':0.001,
  35.     'weight_decay':0.0001
  36. }
  38. transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
  39. mnist_train = MNIST(os.getcwd(), train=True, download=True, transform=transform)
  40. mnist_train = DataLoader(mnist_train, batch_size=256,num_workers=8)
  42. model = LitMNIST(hparams)
  43. trainer = Trainer()
  44. trainer.fit(model, mnist_train)