text_plain_output_1.png
from torch import nn

from torch.utils.data import DataLoader, random_split

from torchvision import transforms

from torchvision.datasets import MNIST

import os

import os

import pytorch_lightning as L

import torch

import torch.nn.functional as F

import numpy as np
import pandas as pd
import os

class TinyModel(torch.nn.Module):
    def __init__(self):
        super(TinyModel, self).__init__()
        self.linear1 = torch.nn.Linear(100, 20)
        self.activation1 = torch.nn.ReLU()
        self.linear2 = torch.nn.Linear(20, 5)
        self.activation2 = torch.nn.ReLU()
        self.linear3 = torch.nn.Linear(5, 2)
    def forward(self, x):
        x = self.linear1(x)
        x = self.activation1(x)
        x = self.linear2(x)
        x = self.activation2(x)
        x = self.linear3(x)
        return x

class LitAutoEncoder(L.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))
    def forward(self, x):
        embedding = self.encoder(x)
        return embedding
    def training_step(self, batch, batch_idx):
        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)
        return loss
    def validation_step(self, batch, batch_idx):
        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, prog_bar=True)
        return loss
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=0.001)
        return optimizer

dataset = MNIST(os.getcwd(), download=True, transform=transforms.ToTensor())
train, val = random_split(dataset, [55000, 5000])
train_dataloader = DataLoader(train, batch_size=64, num_workers=4)
val_dataloader = DataLoader(val, batch_size=64, num_workers=4)
autoencoder = LitAutoEncoder()
trainer = L.Trainer(val_check_interval=250, accelerator='gpu', devices=2, max_epochs=20, precision=16)
trainer.fit(autoencoder, train_dataloader, val_dataloader)