mirror of
https://github.com/webclinic017/drift.git
synced 2026-07-28 11:17:47 +00:00
Daniel Szemerey
d4676e099b
* feat: Refractored and created new model. Pipeline not ready yet. * feat: Implemented and refactored a data pipeline. * ref: Refractored to make more sense. * feat: Training works now with models that you can change. * feat: Added predict function but without working instructions. * feat: gitignore.
60 lines
1.8 KiB
Python
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 |