Files
drift/lightning/model_lightning.py
T
Mark Aron Szulyovszky 797791d2f4 feat(Forecasting): pytorch-forecasting scaffolding is now working, added narrow data format, fixed missing time column index name (#5)
* feat: Started implementing pytorch-forecasting.

* feat(Forecasting): pytorch-forecasting scaffolding is now working, added narrow data format, fixed missing `time` column index name

Co-authored-by: Daniel Szemerey <szemy2@gmail.com>
2021-11-16 10:00:26 +01:00

60 lines
1.8 KiB
Python

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader, random_split
import pytorch_lightning as pl
class LitManualAutoEncoder(pl.LightningModule):
def __init__(self):
super().__init__()
self.encoder = nn.Sequential(nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 3))
self.decoder = nn.Sequential(nn.Linear(3, 128), nn.ReLU(), nn.Linear(128, 28 * 28))
print("success")
def training_step(self, batch, batch_idx):
# --------------------------
# REPLACE WITH YOUR OWN
opt_a = self.optimizers()
x, y = batch
x = x.view(x.size(0), -1)
z = self.encoder(x)
x_hat = self.decoder(z)
loss = F.mse_loss(x_hat, x)
# backward acts like normal backward
self.manual_backward(loss, opt_a, retain_graph=True)
self.manual_backward(loss, opt_a)
opt_a.step()
opt_a.zero_grad()
# --------------------------
def validation_step(self, batch, batch_idx):
# --------------------------
# REPLACE WITH YOUR OWN
x, y = batch
x = x.view(x.size(0), -1)
z = self.encoder(x)
x_hat = self.decoder(z)
loss = F.mse_loss(x_hat, x)
self.log('val_loss', loss)
# --------------------------
def test_step(self, batch, batch_idx):
# --------------------------
# REPLACE WITH YOUR OWN
x, y = batch
x = x.view(x.size(0), -1)
z = self.encoder(x)
x_hat = self.decoder(z)
loss = F.mse_loss(x_hat, x)
self.log('test_loss', loss)
# --------------------------
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
return optimizer