【pytorchで深層生成モデル#5】pytorchでVAE

# [1]
!nvidia-smi

# [2]
# データ作成に使用するライブラリ
from torchvision import datasets
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
# モデル作成に使用するライブラリ
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
# よく使用するライブラリ
import matplotlib.pyplot as plt
import numpy as np
torch.manual_seed(1)

# [3]
# データの読み込み
transform = transforms.Compose([transforms.ToTensor(), lambda x: x.view(-1)])
root = './data'
mnist_dataset = datasets.MNIST(root=root,download=True,train=True,transform=transform)
dataloader = DataLoader(mnist_dataset, batch_size=100,shuffle=True)

# [4]
# gpuの指定
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device

# [1]

!nvidia-smi

​

# [2]

# データ作成に使用するライブラリ

from torchvision import datasets

import torchvision.transforms as transforms

from torch.utils.data import DataLoader

# モデル作成に使用するライブラリ

import torch

import torch.nn as nn

import torch.optim as optim

import torch.nn.functional as F

# よく使用するライブラリ

import matplotlib.pyplot as plt

import numpy as np

torch.manual_seed(1)

​

# [3]

# データの読み込み

transform = transforms.Compose([transforms.ToTensor(), lambda x: x.view(-1)])

root = './data'

mnist_dataset = datasets.MNIST(root=root,download=True,train=True,transform=transform)

dataloader = DataLoader(mnist_dataset, batch_size=100,shuffle=True)

​

# [4]

# gpuの指定

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

device

# [5]
class VAE(nn.Module):
    def __init__(self, device='cpu'):
        super().__init__()
        self.device = device
        self.encoder = Encoder(device=device)
        self.decoder = Decoder(device=device)

    def forward(self, x):
        mean, var = self.encoder(x)
        z = self.reparameterize(mean, var)
        y = self.decoder(z)
        return y, z

    def reparameterize(self, mean, var):
        z = mean + torch.sqrt(var) * torch.randn(mean.size()).to(self.device)
        return z

    def criterion(self, x):
        mean, var = self.encoder(x)
        z = self.reparameterize(mean, var)
        y = self.decoder(z)
        L1 =  - torch.mean(torch.sum(x * torch.log(y) + (1 - x) * torch.log(1 - y), dim=1))
        L2 = - 1/2 * torch.mean(torch.sum(1 + torch.log(var) - mean**2 - var, dim=1))
        L =  L1 + L2

        return L

# [6]
class Encoder(nn.Module):
    def __init__(self, device='cpu'):
        super().__init__()
        self.device = device
        self.l1 = nn.Linear(784, 256)
        self.l2 = nn.Linear(256, 128)
        self.l_mean = nn.Linear(128, 2)
        self.l_var = nn.Linear(128, 2)

    def forward(self, x):
        h = self.l1(x)
        h = torch.relu(h)
        h = self.l2(h)
        h = torch.relu(h)
        mean = self.l_mean(h)
        var = self.l_var(h)
        var = F.softplus(var)

        return mean, var

# [7]
class Decoder(nn.Module):
    def __init__(self, device='cpu'):
        super().__init__()
        self.device = device
        self.l1 = nn.Linear(2, 128)
        self.l2 = nn.Linear(128, 256)
        self.out = nn.Linear(256, 784)

    def forward(self, x):
        h = self.l1(x)
        h = torch.relu(h)
        h = self.l2(h)
        h = torch.relu(h)
        h = self.out(h)
        y = torch.sigmoid(h)

        return y

# [8]
model = VAE(device=device).to(device)
criterion = model.criterion
optimizer = optim.Adam(model.parameters())

# [5]

class VAE(nn.Module):

    def __init__(self, device='cpu'):

        super().__init__()

        self.device = device

        self.encoder = Encoder(device=device)

        self.decoder = Decoder(device=device)

​

    def forward(self, x):

        mean, var = self.encoder(x)

        z = self.reparameterize(mean, var)

        y = self.decoder(z)

        return y, z

​

    def reparameterize(self, mean, var):

        z = mean + torch.sqrt(var) * torch.randn(mean.size()).to(self.device)

        return z

​

    def criterion(self, x):

        mean, var = self.encoder(x)

        z = self.reparameterize(mean, var)

        y = self.decoder(z)

        L1 =  - torch.mean(torch.sum(x * torch.log(y) + (1 - x) * torch.log(1 - y), dim=1))

        L2 = - 1/2 * torch.mean(torch.sum(1 + torch.log(var) - mean**2 - var, dim=1))

        L =  L1 + L2

​

        return L

​

# [6]

class Encoder(nn.Module):

    def __init__(self, device='cpu'):

        super().__init__()

        self.device = device

        self.l1 = nn.Linear(784, 256)

        self.l2 = nn.Linear(256, 128)

        self.l_mean = nn.Linear(128, 2)

        self.l_var = nn.Linear(128, 2)

​

    def forward(self, x):

        h = self.l1(x)

        h = torch.relu(h)

        h = self.l2(h)

        h = torch.relu(h)

        mean = self.l_mean(h)

        var = self.l_var(h)

        var = F.softplus(var)

​

        return mean, var

​

# [7]

class Decoder(nn.Module):

    def __init__(self, device='cpu'):

        super().__init__()

        self.device = device

        self.l1 = nn.Linear(2, 128)

        self.l2 = nn.Linear(128, 256)

        self.out = nn.Linear(256, 784)

​

    def forward(self, x):

        h = self.l1(x)

        h = torch.relu(h)

        h = self.l2(h)

        h = torch.relu(h)

        h = self.out(h)

        y = torch.sigmoid(h)

​

        return y

​

# [8]

model = VAE(device=device).to(device)

criterion = model.criterion

optimizer = optim.Adam(model.parameters())

# [9]
n_epoch = 8

for epoch in range(n_epoch):
    loss_mean = 0.
    for (x, t) in dataloader:

      # 学習準備
      x = x.to(device)
      model.train()
      
      # モデルの学習
      loss = criterion(x)
      optimizer.zero_grad()
      loss.backward()     
      optimizer.step()

      # 損失関数の計算
      loss_mean += loss.item()
    loss_mean /= len(dataloader)

    print('Epoch: {}, Loss: {:.3f}'.format(epoch+1, loss_mean))

# [9]

n_epoch = 8

​

for epoch in range(n_epoch):

    loss_mean = 0.

    for (x, t) in dataloader:

​

      # 学習準備

      x = x.to(device)

      model.train()

      # モデルの学習

      loss = criterion(x)

      optimizer.zero_grad()

      loss.backward()     

      optimizer.step()

​

      # 損失関数の計算

      loss_mean += loss.item()

    loss_mean /= len(dataloader)

​

    print('Epoch: {}, Loss: {:.3f}'.format(epoch+1, loss_mean))

# [10]
# データの生成
model.eval()
z = torch.randn(10, 2, device = device)
images = model.decoder(z)
images = images.view(-1, 28, 28)
images = images.squeeze().detach().cpu().numpy()

# データの可視化
for i, image in enumerate(images):
    plt.subplot(2, 5, i+1)
    plt.imshow(image, cmap='binary_r')
    plt.axis('off')
plt.tight_layout()
plt.show()

# [11]
# データ可視化の前準備
img_size=28
n_image = 10
img_size_spaced = img_size + 2
matrix_image = np.zeros((img_size_spaced*n_image, img_size_spaced*n_image))  # 全体の画像

# 潜在変数の作成
z_1 = torch.linspace(-3, 3, n_image)  # 行
z_2 = torch.linspace(-3, 3, n_image)  # 列

#  潜在変数を変化させて画像を生成
for i, z1 in enumerate(z_1):
    for j, z2 in enumerate(z_2):
        x = torch.tensor([float(z1), float(z2)], device=device)
        images = model.decoder(x)
        images = images.view(-1, 28, 28)
        images = images.squeeze().detach().cpu().numpy()
        top = i*img_size_spaced
        left = j*img_size_spaced
        matrix_image[top : top+img_size, left : left+img_size] = images

# データの可視化
plt.figure(figsize=(8, 8))
plt.imshow(matrix_image.tolist(), cmap="Greys_r")
plt.tick_params(labelbottom=False, labelleft=False, bottom=False, left=False)
plt.show()

# [10]

# データの生成

model.eval()

z = torch.randn(10, 2, device = device)

images = model.decoder(z)

images = images.view(-1, 28, 28)

images = images.squeeze().detach().cpu().numpy()

​

# データの可視化

for i, image in enumerate(images):

    plt.subplot(2, 5, i+1)

    plt.imshow(image, cmap='binary_r')

    plt.axis('off')

plt.tight_layout()

plt.show()

​

# [11]

# データ可視化の前準備

img_size=28

n_image = 10

img_size_spaced = img_size + 2

matrix_image = np.zeros((img_size_spaced*n_image, img_size_spaced*n_image))  # 全体の画像

​

# 潜在変数の作成

z_1 = torch.linspace(-3, 3, n_image)  # 行

z_2 = torch.linspace(-3, 3, n_image)  # 列

​

#  潜在変数を変化させて画像を生成

for i, z1 in enumerate(z_1):

    for j, z2 in enumerate(z_2):

        x = torch.tensor([float(z1), float(z2)], device=device)

        images = model.decoder(x)

        images = images.view(-1, 28, 28)

        images = images.squeeze().detach().cpu().numpy()

        top = i*img_size_spaced

        left = j*img_size_spaced

        matrix_image[top : top+img_size, left : left+img_size] = images

​

# データの可視化

plt.figure(figsize=(8, 8))

plt.imshow(matrix_image.tolist(), cmap="Greys_r")

plt.tick_params(labelbottom=False, labelleft=False, bottom=False, left=False)

plt.show()