Files
pytorch-lightning/tests/base/debug.py
T
Adrian WälchliandJirka Borovec b7de42f70d Add MNIST dataset & drop torchvision dep. from tests (#986)
* added custom mnist without torchvision dep

* move files so it does not conflict with mnist gitignore

* mock torchvision for tests

* fix line too long

* fix line too long

* fix "module level import not at top of file" warning

* move mock imports to __init__.py

* simplify MNIST a lot and download directly the .pt files

* further simplify and clean up mnist

* revert import overrides

* make as before

* drop  PIL requirement

* move mnist.py to datasets subfolder

* use logging instead of print

* choose same name as in torchvision

* remove torchvision and pillow also from yml file

* refactor if train

Co-Authored-By: Jirka Borovec <Borda@users.noreply.github.com>

* capitalized class attr

* moved mnist to models

* re-added datsets ignore

* better name for file variable

* Update mnist.py

* move dataset classes to datasets.py

* new line

* update

* update

* fix automerge

* move to base folder

* adapt testingmnist to new mnist base class

* remove temporal fix

* fix datatype

* remove old testingmnist

* readable

* fix import

* fix whitespace

* docstring

Co-Authored-By: Jirka Borovec <Borda@users.noreply.github.com>

* Update tests/base/datasets.py

Co-Authored-By: Jirka Borovec <Borda@users.noreply.github.com>

* changelog

* added types

* Update CHANGELOG.md

Co-Authored-By: Jirka Borovec <Borda@users.noreply.github.com>

* exist->isfile

Co-Authored-By: Jirka Borovec <Borda@users.noreply.github.com>

* index -> idx

* temporary fix for trains error

* better changelog message

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
2020-03-30 18:25:37 -04:00

52 lines
1.5 KiB
Python

import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader
import pytorch_lightning as pl
from tests.base.datasets import MNIST
# from test_models import assert_ok_test_acc, load_model, \
# clear_save_dir, get_default_testtube_logger, get_default_hparams, init_save_dir, \
# init_checkpoint_callback, reset_seed, set_random_master_port
class CoolModel(pl.LightningModule):
def __init(self):
super().__init__()
# not the best model...
self.l1 = torch.nn.Linear(28 * 28, 10)
def forward(self, x):
return torch.relu(self.l1(x))
def my_loss(self, y_hat, y):
return F.cross_entropy(y_hat, y)
def training_step(self, batch, batch_idx):
x, y = batch
y_hat = self(x)
return {'training_loss': self.my_loss(y_hat, y)}
def validation_step(self, batch, batch_idx):
x, y = batch
y_hat = self(x)
return {'val_loss': self.my_loss(y_hat, y)}
def validation_epoch_end(self, outputs):
avg_loss = torch.stack([x for x in outputs['val_loss']]).mean()
return avg_loss
def configure_optimizers(self):
return [torch.optim.Adam(self.parameters(), lr=0.02)]
def train_dataloader(self):
return DataLoader(MNIST('path/to/save', train=True), batch_size=32)
def val_dataloader(self):
return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
def test_dataloader(self):
return DataLoader(MNIST('path/to/save', train=False), batch_size=32)