DL_Final_Project/vae_example.py at main · mozesjacobs/DL_Final_Project · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
# -*- coding: utf-8 -*-
"""VAE_example.ipynb

Automatically generated by Colaboratory.

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

import numpy as np
import torch
import torch.nn as nn

class VAE(nn.Module):
    def __init__(self, device, batch_size=250):
        super(VAE, self).__init__()

        self.device = device

        self.encoder = nn.Sequential(
            nn.Conv2d(3, 32, 4, stride=2),
            nn.ReLU(),
            nn.Conv2d(32, 64, 4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 128, 4, stride=2),
            nn.ReLU(),
            nn.Conv2d(128, 256, 4, stride=2),
            nn.ReLU()
        )

        self.mufc = nn.Linear(1024, 32)
        self.logvarfc = nn.Linear(1024, 32)

        self.decoder_fc = nn.Linear(32, 1024)

        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(1024, 128, 5, stride=2),
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, 5, stride=2),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, 6, stride=2),
            nn.ReLU(),
            nn.ConvTranspose2d(32, 3, 6, stride=2),
            nn.Sigmoid(),
        )

        self.batch_size = batch_size
        #self.dist = torch.distributions.laplace.Laplace(0, torch.ones([50]))

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        #noise = torch.randn(self.batch_size, 32).to(self.device)
        noise = torch.randn_like(std).to(self.device)
        return mu + std * noise # z

    def forward(self, x):
        x = self.encoder(x)
        x = x.reshape(-1, 1024)
        mu, logvar = self.mufc(x), self.logvarfc(x)
        z = self.reparameterize(mu, logvar)
        z_ = self.decoder_fc(z)
        z_ = z_.reshape(-1, 1024, 1, 1)
        return self.decoder(z_.float()), mu, logvar

    def get_z(self, x):
        with torch.no_grad():
            encoded = self.encoder(x)
            return self.reparameterize(encoded)

    def loss_func(self, x, x_prime, mu, logvar):
      recon_loss = nn.BCELoss(reduction='sum')
      loss = recon_loss(x_prime, x)
      loss += -0.5 * torch.sum(1 + logvar - mu.pow(2) - torch.exp(logvar))

      return loss

USE_CUDA = True
use_cuda = USE_CUDA and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
print('Using device', device)
import multiprocessing
NUM_WORKERS = multiprocessing.cpu_count()
print('num workers:', NUM_WORKERS)

vae = VAE(device, 250)
vae.to(device, dtype=torch.float)

from google.colab import drive
drive.mount('/content/drive')

vae_path = '/content/drive/MyDrive/vae_original.pt'
vae.load_state_dict(torch.load(vae_path))