text_plain_output_5.png application_vnd.jupyter.stderr_output_2.png application_vnd.jupyter.stderr_output_4.png text_plain_output_3.png text_plain_output_1.png
from IPython.display import clear_output

from skimage import io

from sklearn.model_selection import train_test_split

from time import time

from torch.optim import lr_scheduler

from torchvision import datasets, transforms

from tqdm.autonotebook import tqdm

import matplotlib.pyplot as plt

import numpy as np

import os

import pandas as pd

import torch

import torch.nn as nn

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

def fetch_dataset(path, attrs_name='lfw_attributes.txt', images_name='lfw-deepfunneled', dx=80, dy=80, dimx=64, dimy=64):
    if not os.path.exists(images_name):
        os.system('wget http://vis-www.cs.umass.edu/lfw/lfw-deepfunneled.tgz -O tmp.tgz')
        os.system('tar xvzf tmp.tgz && rm tmp.tgz')
        assert os.path.exists(images_name)
    if not os.path.exists(attrs_name):
        os.system('wget http://www.cs.columbia.edu/CAVE/databases/pubfig/download/%s' % attrs_name)
    transform = transforms.Compose([transforms.ToPILImage(), transforms.CenterCrop((dx, dy)), transforms.Resize((dimx, dimy)), transforms.ToTensor()])
    df_attrs = pd.read_csv(os.path.join(path, attrs_name), sep='\t', skiprows=1)
    df_attrs = pd.DataFrame(df_attrs.iloc[:, :-1].values, columns=df_attrs.columns[1:])
    photo_ids = []
    for dirpath, dirnames, filenames in os.walk(os.path.join(path, images_name)):
        for fname in filenames:
            if fname.endswith('.jpg'):
                fpath = os.path.join(dirpath, fname)
                photo_id = fname[:-4].replace('_', ' ').split()
                person_id = ' '.join(photo_id[:-1])
                photo_number = int(photo_id[-1])
                photo_ids.append({'person': person_id, 'imagenum': photo_number, 'photo_path': fpath})
    photo_ids = pd.DataFrame(photo_ids)
    df = pd.merge(df_attrs, photo_ids, on=('person', 'imagenum'))
    assert len(df) == len(df_attrs), 'lost some data when merging dataframes'
    all_photos = df['photo_path'].apply(io.imread).apply(transform)
    all_photos = all_photos.values
    all_attrs = df.drop(['photo_path', 'person', 'imagenum'], axis=1)
    return (all_photos, all_attrs)

batch_size = 32
train_photos, val_photos, train_attrs, val_attrs = train_test_split(data, attrs, train_size=0.8, shuffle=False)
data_tr = torch.utils.data.DataLoader(train_photos, batch_size=batch_size)
data_val = torch.utils.data.DataLoader(val_photos, batch_size=batch_size)

for i in range(12):
    plt.axis('off')

dim_code = 32

class CVAE(nn.Module):
    def __init__(self, base_channel_size: int, latent_dim: int, num_classes: int, num_input_channels: int=3, act_fn=nn.ReLU):
        super().__init__()
        self.dummy_param = nn.Parameter(torch.empty(0))
        self.latent_dim = latent_dim
        self.c_hid = base_channel_size
        self.num_classes = num_classes
        conv_size = int(np.exp2(np.log2(self.c_hid) - 3))
        ln_size = 2 * self.c_hid * conv_size * conv_size
        self.encoder = nn.Sequential(nn.Conv2d(num_input_channels, self.c_hid, kernel_size=3, padding=1, stride=2), act_fn(), nn.Conv2d(self.c_hid, self.c_hid, kernel_size=3, padding=1), act_fn(), nn.Conv2d(self.c_hid, 2 * self.c_hid, kernel_size=3, padding=1, stride=2), act_fn(), nn.Conv2d(2 * self.c_hid, 2 * self.c_hid, kernel_size=3, padding=1), act_fn(), nn.Conv2d(2 * self.c_hid, 2 * self.c_hid, kernel_size=3, padding=1, stride=2), act_fn())
        self.flatten = nn.Flatten(start_dim=1)
        self.linear_mu = nn.Sequential(nn.Linear(ln_size, latent_dim))
        self.linear_logvar = nn.Sequential(nn.Linear(ln_size, latent_dim))
        self.linear_decoder = nn.Sequential(nn.Linear(latent_dim + num_classes, ln_size), act_fn())
        self.unflatten = nn.Sequential(nn.Unflatten(dim=1, unflattened_size=(2 * self.c_hid, conv_size, conv_size)))
        self.decoder = nn.Sequential(nn.ConvTranspose2d(2 * self.c_hid, 2 * self.c_hid, kernel_size=3, output_padding=1, padding=1, stride=2), act_fn(), nn.Conv2d(2 * self.c_hid, 2 * self.c_hid, kernel_size=3, padding=1), act_fn(), nn.ConvTranspose2d(2 * self.c_hid, self.c_hid, kernel_size=3, output_padding=1, padding=1, stride=2), act_fn(), nn.Conv2d(self.c_hid, self.c_hid, kernel_size=3, padding=1), act_fn(), nn.ConvTranspose2d(self.c_hid, num_input_channels, kernel_size=3, output_padding=1, padding=1, stride=2), nn.Sigmoid())
    def encode(self, x):
        x = self.encoder(x)
        x = self.flatten(x)
        mu = self.linear_mu(x)
        logvar = self.linear_logvar(x)
        return (mu, logvar)
    def reparameterize(self, mu, logvar):
        if self.training:
            std = torch.exp(logvar / 2)
            eps = torch.randn_like(std)
            return eps * std + mu
        else:
            return mu
    def decode(self, x):
        x = self.linear_decoder(x)
        x = self.unflatten(x)
        x = self.decoder(x)
        return x
    def forward(self, x, **kwargs):
        y = kwargs['labels']
        y = torch.nn.functional.one_hot(y, num_classes=self.num_classes)
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        z = torch.cat([z, y], dim=1)
        z = self.decode(z)
        return (mu, logvar, z)
    def sample(self, labels: list):
        y = torch.tensor(labels, dtype=torch.int64).to(self.dummy_param.device)
        y = torch.nn.functional.one_hot(y, num_classes=self.num_classes)
        z = torch.randn(y.size()[0], 32).to(self.dummy_param.device)
        z = torch.cat([z, y], dim=1)
        return self.decode(z)

def loss_vae(x, mu, logsigma, reconstruction):
    kl = KL_divergence(mu, logsigma)
    ll = log_likelihood(x, reconstruction)
    return kl + ll

def KL_divergence(mu, logvar):
    """
    часть функции потерь, которая отвечает за "близость" латентных представлений разных людей
    """
    loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return loss
def log_likelihood(x, reconstruction):
    """
    часть функции потерь, которая отвечает за качество реконструкции (как mse в обычном autoencoder)
    """
    loss = nn.BCELoss(reduction='sum')
    return loss(reconstruction, x)
def loss_vae(x, mu, logsigma, reconstruction):
    kl = KL_divergence(mu, logsigma)
    ll = log_likelihood(x, reconstruction)
    return kl + ll

batch_size = 32
size = 32
transform = transforms.Compose([transforms.Resize(size), transforms.ToTensor()])
train_dataset = datasets.MNIST(root='./mnist_data/', transform=transform, train=True, download=True)
test_dataset = datasets.MNIST(root='./mnist_data/', transform=transform, train=False, download=False)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

criterion = loss_vae
autoencoder = CVAE(num_input_channels=1, base_channel_size=32, num_classes=train_dataset.targets.unique().size()[0], latent_dim=dim_code)
optimizer = torch.optim.Adam(autoencoder.parameters(), lr=0.001)

def train(model, opt, loss_fn, epochs, data_tr, data_val, scheduler=None, device='cpu', show=True, show_num=3):
    from time import time
    from tqdm.autonotebook import tqdm
    model = model.to(device)
    X_val, Y_val = next(iter(data_val))
    train_losses = []
    val_losses = []
    log_template = 'Epoch {ep:03d}/{epochs:03d}  train loss: {t_loss:0.4f} val loss {v_loss:0.4f}'
    with tqdm(desc='epoch', total=epochs) as pbar_outer:
        for epoch in range(epochs):
            tic = time()
            avg_loss = 0
            model.train()
            for X_batch, Y_batch in data_tr:
                X_batch = X_batch.to(device, dtype=torch.float32)
                Y_batch = Y_batch.to(device)
                opt.zero_grad()
                mu, logvar, X_pred = model(X_batch, labels=Y_batch)
                loss = loss_fn(X_batch, mu, logvar, X_pred)
                loss.backward()
                opt.step()
                avg_loss += loss / len(data_tr)
            toc = time()
            model.eval()
            mu, logvar, X_hat = model(X_val.to(device, dtype=torch.float32), labels=Y_val.to(device))
            X_hat = X_hat.detach().to('cpu')
            train_losses.append(avg_loss.item())
            val_losses.append(loss_fn(X_val, mu, logvar, X_hat).item())
            nums = np.random.randint(10, size=show_num)
            output_nums = model.sample(nums).detach()
            output_nums = output_nums.detach().to('cpu')
            if scheduler:
                scheduler.step()
            pbar_outer.update(1)
            if show:
                clear_output(wait=True)
                plt.clf()
                for k in range(show_num):
                    plt.axis('off')
                    plt.axis('off')
                    plt.axis('off')
            else:
                tqdm.write(log_template.format(ep=epoch + 1, epochs=epochs, t_loss=train_losses[-1], v_loss=val_losses[-1]))
    return (train_losses, val_losses, X_hat, mu, logvar)

from torch.optim import lr_scheduler
lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

max_epochs = 20
cvae_train_loss, cvae_val_loss, cvae_predict_img_val, cvae_mu, cvae_logvar = train(model=autoencoder, opt=optimizer, loss_fn=criterion, epochs=max_epochs, data_tr=train_loader, data_val=test_loader, device=device, scheduler=lr_scheduler, show=True)