text_plain_output_1.png
from skimage import io

from sklearn.model_selection import train_test_split

from torchvision import datasets, transforms

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)

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)