vae.py

# -*- coding: utf-8 -*-
"""vae

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1UZXVNVS2vlj4zyeOs_kh2nsvX9zu6ioz
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

class VAE(nn.Module):
    def __init__(self, input_channels, input_height, input_width, latent_dim):
        super(VAE, self).__init__()

        # Encoder
        self.encoder = nn.Sequential(
            nn.Conv2d(input_channels, 32, kernel_size=3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(64 * (input_height // 4) * (input_width // 4), 256),
            nn.ReLU()
        )

        # Temporal encoding with LSTM
        self.lstm = nn.LSTM(256, hidden_size=128, num_layers=1, batch_first=True)

        # Sampling layer
        self.sampling = self._sampling

        # Decoder
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 64 * (input_height // 4) * (input_width // 4)),
            nn.ReLU(),
            nn.Unflatten(1, (64, input_height // 4, input_width // 4)),
            nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(32, input_channels, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.Sigmoid()
        )

    def _sampling(self, mean, log_var):
        std = torch.exp(0.5 * log_var)
        eps = torch.randn_like(std)
        return mean + eps * std

    def forward(self, x):

        x = self.encoder(x)


        x = x.unsqueeze(1)
        x, _ = self.lstm(x)
        x = x[:, -1, :]

        mean, log_var = x.chunk(2, dim=1)


        z = self.sampling(mean, log_var)


        x_recon = self.decoder(z)

        return x_recon, mean, log_var


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


latent_dim = 128
input_channels = 3
input_height = 64
input_width = 64


vae = VAE(input_channels, input_height, input_width, latent_dim).to(device)

print(vae)


criterion = nn.MSELoss()
optimizer = optim.Adam(vae.parameters(), lr=0.001)

# Training loop
for epoch in range(num_epochs):
    for data in train_loader:
        inputs = data.to(device)


        outputs, mean, log_var = vae(inputs)


        reconstruction_loss = criterion(outputs, inputs)
        kl_divergence = -0.5 * torch.sum(1 + log_var - mean.pow(2) - log_var.exp())


        loss = reconstruction_loss + kl_divergence


        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (i+1) % 100 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{total_steps}], Loss: {loss.item():.4f}')


    vae.eval()
    val_loss = 0.0
    with torch.no_grad():
        for val_data in val_loader:
            val_inputs = val_data.to(device)
            val_outputs, val_mean, val_log_var = vae(val_inputs)
            val_loss += criterion(val_outputs, val_inputs).item()

    val_loss /= len(val_loader)

    if val_loss < best_val_loss:
        best_val_loss = val_loss
        torch.save(vae.state_dict(), 'best_model.pth')

    print(f'Epoch [{epoch+1}/{num_epochs}], Validation Loss: {val_loss:.4f}')

transform = torchvision.transforms.ToTensor()  # Adjust based on your preprocessing needs
dataset = VideoDataset(root_dir='path/to/your/data', transform=transform)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)