| 32 | |
| 33 | |
| 34 | class LitClassifier(pl.LightningModule): |
| 35 | |
| 36 | def __init__(self): |
| 37 | super().__init__() |
| 38 | self.l1 = torch.nn.Linear(28 * 28, 10) |
| 39 | |
| 40 | def forward(self, x): |
| 41 | return F.relu(self.l1(x.view(x.size(0), -1))) |
| 42 | |
| 43 | def training_step(self, batch, batch_idx): |
| 44 | x, y = batch |
| 45 | y_hat = self(x) |
| 46 | loss = F.cross_entropy(y_hat, y) |
| 47 | self.log('train_loss', loss) |
| 48 | return loss |
| 49 | |
| 50 | def validation_step(self, batch, batch_idx): |
| 51 | x, y = batch |
| 52 | y_hat = self(x) |
| 53 | loss = F.cross_entropy(y_hat, y) |
| 54 | self.log('val_loss', loss) |
| 55 | |
| 56 | def configure_optimizers(self): |
| 57 | return torch.optim.Adam(self.parameters(), lr=1e-2) |
| 58 | |
| 59 | class TestPytorchLightning(unittest.TestCase): |
| 60 | |