<kaggle_start><data_title>Flickr-Faces-HQ (FFHQ) small<data_description>This is a small version of FFHQ with 3143 photos. 
Flickr-Faces-HQ (FFHQ) is an image dataset containing high-quality images of human faces. It is provided by NVIDIA under the Creative Commons BY-NC-SA 4.0 license. It offers 70,000 PNG images at 1024×1024 resolution that display diverse ages, ethnicities, image backgrounds, and accessories like hats and eyeglasses.<data_name>faces-dataset-small
<code>#
# Школа глубокого обучения ФПМИ МФТИ
# Домашнее задание. Generative adversarial networks
# В этом домашнем задании вы обучите GAN генерировать лица людей и посмотрите на то, как можно оценивать качество генерации
import os
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
import torchvision.transforms as tt
import torchvision
import torch
import torch.nn as nn
import cv2
from tqdm.notebook import tqdm
from torchvision.utils import save_image
from torchvision.utils import make_grid
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from torch.utils.data import Dataset
import multiprocessing as mp
from PIL import Image
from IPython.display import clear_output

sns.set(style="darkgrid", font_scale=1.2)
sns.set_style("darkgrid")


# ## Часть 1. Подготовка данных (1 балл)
# В качестве обучающей выборки возьмем часть датасета [Flickr Faces](https://github.com/NVlabs/ffhq-dataset), который содержит изображения лиц людей в высоком разрешении (1024х1024). Оригинальный датасет очень большой, поэтому мы возьмем его часть. Скачать датасет можно [здесь](https://www.kaggle.com/datasets/tommykamaz/faces-dataset-small?resource=download-directory) и  [здесь](https://drive.google.com/drive/folders/14H7LQWzd09SaM11oLfGtd455sWH2rmRW?usp=share_link)
# Давайте загрузим наши изображения. Напишите функцию, которая строит DataLoader для изображений, при этом меняя их размер до нужного значения (размер 1024 слишком большой, поэтому мы рекомендуем взять размер 128 либо немного больше)
class FacesDataset(Dataset):
    """
    Dataset for Flickr-Faces-HQ (FFHQ) small
    """

    def __init__(self, path, input_transform=None):
        super().__init__()
        self.path = path
        self.input_transform = input_transform
        self.paths = os.listdir(self.path)
        self.data = []
        self._read_images()

    def _read_image(self, image_name):
        """
        Read image and use input_transforms
        """
        image = Image.open(os.path.join(self.path, image_name))
        if self.input_transform:
            image = self.input_transform(image)
        return image

    def _read_images(self):
        """
        Pool of _read_image functions
        """
        with mp.Pool(processes=mp.cpu_count()) as pool:
            self.data = list(
                tqdm(
                    pool.imap_unordered(self._read_image, self.paths),
                    total=len(self.paths),
                )
            )

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]


def to_device(data, device):
    """
    Move data to self.device
    """
    if isinstance(data, (list, tuple)):
        return [to_device(x, device) for x in data]
    return data.to(device, non_blocking=True)


class DeviceDataLoader:
    def __init__(self, dataloader, device):
        self.dataloader = dataloader
        self.device = device

    def __iter__(self):
        for batch in self.dataloader:
            yield to_device(batch, self.device)

    def __len__(self):
        return len(self.dataloader)


def get_dataloader(path, device, input_transform, image_size, batch_size):
    """
    Builds dataloader for training data.
    Use tt.Compose and tt.Resize for transformations
    :param image_size: height and wdith of the image
    :param batch_size: batch_size of the dataloader
    :returns: DataLoader object
    """
    # TODO: resize images, convert them to tensors and build dataloader
    faces_dataset = FacesDataset(path, input_transform)
    return DeviceDataLoader(
        DataLoader(faces_dataset, batch_size=batch_size, drop_last=True), device
    )


# !ls /kaggle/input/faces-dataset-small/faces_dataset_small
image_size = (128, 128)
batch_size = 64
path = "/kaggle/input/faces-dataset-small/faces_dataset_small"
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
input_transform = tt.Compose(
    [
        tt.RandomHorizontalFlip(p=0.5),
        tt.Resize(image_size),
        tt.ToTensor(),
        tt.ConvertImageDtype(torch.float),
    ]
)
faces_dataloader = get_dataloader(path, device, input_transform, image_size, batch_size)
faces = next(iter(faces_dataloader))
grid = torchvision.utils.make_grid(faces, nrow=8, padding=0, scale_each=True)
fig = plt.figure(figsize=(16, 8))
plt.imshow(grid.cpu().permute(1, 2, 0))
plt.axis("off")
plt.show()


# ## Часть 2. Построение и обучение модели (2 балла)
# Сконструируйте генератор и дискриминатор. Помните, что:
# * дискриминатор принимает на вход изображение (тензор размера `3 x image_size x image_size`) и выдает вероятность того, что изображение настоящее (тензор размера 1)
# * генератор принимает на вход тензор шумов размера `latent_size x 1 x 1` и генерирует изображение размера `3 x image_size x image_size`
# ### PixelNorm
# to avoid training BatchNorm module
class PixelNorm(nn.Module):
    """
    Nontrainable pixel normalization layer
    """

    def __init__(self, eps=1e-8):
        super(PixelNorm, self).__init__()
        self.eps = eps

    def forward(self, x):
        return (
            x
            / torch.sqrt(torch.mean(x**2, dim=1, keepdim=True) + self.eps)
            / x.shape[0]
        )


# Что же такое PixelNorm?
# В [статье по ProGAN]("https://towardsdatascience.com/progan-how-nvidia-generated-images-of-unprecedented-quality-51c98ec2cbd2") предлагается использовать такую нормализация вместо BatchNorm. Главным отличием от BatchNorm является отсутствие тренировочных весов у этого слоя.
# $$b_{x,y}=\frac{a_{x,y}}{\sqrt{\frac{1}{C}\sum\limits_{j=0}^C (a_{x,y}^j)^2 + \varepsilon}}$$
# Но как же влияет PixelNorm на фотографию?
faces = next(iter(faces_dataloader))
grid = torchvision.utils.make_grid(faces, nrow=8, padding=0, scale_each=True)
fig = plt.figure(figsize=(16, 8))
# Using PixelNorm
plt.imshow(PixelNorm()(grid).cpu().permute(1, 2, 0))
plt.axis("off")
plt.show()


# Норм, однако давайте не использовать нормализацию на последнем слое генератора, чтобы он смог восстанавливать цвета.
# ### Generator
class DCGenerator(nn.Module):
    """
    Generator of Deep Convolution GAN
    """

    def __init__(self, z_dim=96, channels_image=3, features_g=8):
        super(DCGenerator, self).__init__()
        self.channels_image = channels_image
        self.latent_size = z_dim
        self._block1 = self._block(
            z_dim, features_g * 32, 4, 1, 0
        )  # N x f_g*16 x 4 x 4
        self._block2 = self._block(
            features_g * 32, features_g * 16, 3, 2, 1, output_padding=1
        )  # 8 x 8
        self._block3 = self._block(
            features_g * 16, features_g * 8, 3, 2, 1, output_padding=1
        )  # 16 x 16
        self._block4 = self._block(
            features_g * 8, features_g * 4, 3, 2, 1, output_padding=1
        )  # 32 x 32
        self._block5 = self._block(features_g * 4, features_g * 2, 4, 2, 1)  # 64 x 64
        self.end_block = nn.Sequential(
            nn.ConvTranspose2d(
                features_g * 2, self.channels_image, 4, 2, 1, output_padding=0
            ),  # 128 x 128
            nn.Tanh(),  # [0, +inf] -> [0, 1]
        )

    def _block(
        self, in_channels, out_channels, kernel_size, stride, padding, output_padding=0
    ):
        """
        Generator block consist such elems as
            nn.ConvTranspose2d
            PixelNorm()
            nn.ReLU()
        """
        return nn.Sequential(
            nn.ConvTranspose2d(
                in_channels,
                out_channels,
                kernel_size,
                stride,
                padding,
                output_padding=output_padding,
                bias=False,
            ),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
        )

    def _get_input_tensor_shape(self):
        """
        Return input tensor size
        """
        return self.latent_size

    def forward(self, x):
        x = self._block1(x)
        x = self._block2(x)
        x = tt.GaussianBlur(3)(x)  # Blur Image
        x = self._block3(x)
        # x = tt.GaussianBlur(3)(x) # Blur image
        x = self._block4(x)
        x = self._block5(x)
        x = self.end_block(x)
        return x


# ### Disciminator
class DCDiscriminator(nn.Module):
    """
    Discriminator of Deep Convolution GAN
    """

    def __init__(self, channels_image, features_d, H, W):
        super(DCDiscriminator, self).__init__()
        self.channels_image = channels_image
        self.H = H
        self.W = W
        self.disc = nn.Sequential(
            # N x channels_image x 128 x 128
            nn.Conv2d(
                channels_image, features_d, kernel_size=4, stride=2, padding=1
            ),  # 64 x 64
            nn.LeakyReLU(0.2),
            self._block(features_d, features_d * 2, 4, 2, 1),  # 32 x 32
            self._block(features_d * 2, features_d * 4, 4, 2, 1),  # 16 x 16
            self._block(features_d * 4, features_d * 8, 4, 2, 1),  # 8 x 8
            self._block(features_d * 8, features_d * 16, 4, 2, 1),  # 4 x 4
            nn.Conv2d(features_d * 16, 1, kernel_size=4, stride=2, padding=0),  # 1 x 1
            nn.Sigmoid(),
        )

    def _block(self, in_channels, out_channels, kernel_size, stride, padding):
        """
        Discriminator block consist such elems as
            nn.Conv2d
            nn.BatchNorm2d
            nn.LeakyReLU
        """
        return nn.Sequential(
            nn.Conv2d(
                in_channels, out_channels, kernel_size, stride, padding, bias=False
            ),
            nn.BatchNorm2d(out_channels),
            nn.LeakyReLU(0.1),
        )

    def _get_input_tensor_shape(self):
        """
        Return input tensor size
        """
        return self.channels_image, self.H, self.W

    def forward(self, x):
        return self.disc(x)


# Проинизицализируем веса нормальным распределением с стандартным отлонением 0 и дисперсией 0.02
def initialize_weights(model):
    """
    Model  normal initialization
    Args:
        model: nn.Module object
    """
    for m in model.modules():
        if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.BatchNorm2d)):
            nn.init.normal_(m.weight, 0.0, 0.02)


# Проверим размерности выходов генератора и дискриминатора
def checkout(gen, disc, latent_size, N, in_channels, H, W):
    """
    Checking right shape of generator and discriminator
    Args:
        gen: initilized generator
        disc: initilized discriminator
        latent_size: lenght of latent vector
        N: batch size
        in_channels: source image channels
        (H, W): source image size
    """
    x = torch.randn((N, in_channels, H, W))
    assert disc(x).shape == (N, 1, 1, 1), "Discriminator has wrong output shape"
    z = torch.randn((N, latent_size, 1, 1))
    assert gen(z).shape == (
        N,
        in_channels,
        H,
        W,
    ), "Generator has wrong output shape: {}. \nExpected ({}, {}, {}, {})".format(
        gen(z).shape, N, in_channels, H, W
    )
    print("Success")


latent_size = 96  # choose latent size
batch_size = 8
image_channels = 3
H, W = 128, 128
features = 32
DCdisc = DCDiscriminator(image_channels, features, H, W)
initialize_weights(DCdisc)
DCgen = DCGenerator(latent_size, image_channels, features)
initialize_weights(DCgen)
# do some test
checkout(DCgen, DCdisc, latent_size, batch_size, image_channels, H, W)


# ### Trainer
# Перейдем теперь к обучению нашего GANа. Алгоритм обучения следующий:
# 1. Учим дискриминатор:
# * берем реальные изображения и присваиваем им метку 1
# * генерируем изображения генератором и присваиваем им метку 0
# * обучаем классификатор на два класса
# 2. Учим генератор:
# * генерируем изображения генератором и присваиваем им метку 0
# * предсказываем дискриминаторором, реальное это изображение или нет
# В качестве функции потерь берем бинарную кросс-энтропию
class GANTrainer:
    def __init__(
        self,
        models,
        optimizers,
        criterions,
        latent_size,
        batch_size,
        device,
        plot_step=10,
    ):
        self.device = device
        # init and move to device
        self.model = {name: model.to(self.device) for name, model in models.items()}
        self.criterion = {
            name: criterion.to(self.device) for name, criterion in criterions.items()
        }
        self.optimizer = optimizers
        self.latent_size = latent_size
        self.batch_size = batch_size
        # Plot settings
        self.plot_step = plot_step
        # Losses
        self.losses = {"discriminator": [], "generator": []}
        self.score = {"real": [0], "fake": [0]}

    def train(self, dataloader, epochs):
        """
        Main train GAN function
        """
        self.model["discriminator"].train()
        self.model["generator"].train()
        torch.cuda.empty_cache()
        for epoch in range(epochs):
            self.__prepare_stat()
            print("Epoch {} out of {}".format(epoch + 1, epochs))
            for i, real_image in enumerate(tqdm(dataloader)):
                # Train discriminator
                if i < 5 or self.score["real"][-2] < 0.9:
                    self.optimizer["discriminator"].zero_grad()  # Clear grads
                    loss_d = self.__train_disc(
                        real_image, self._get_fake_images(self.batch_size)
                    )
                    loss_d.backward()  # Calculate gradients
                    self.optimizer["discriminator"].step()  # Doing optimazer step
                # Train Generator
                self.optimizer["generator"].zero_grad()
                loss_g = self.__train_gen(self._get_fake_images(self.batch_size))
                loss_g.backward()
                self.optimizer["generator"].step()
            # Plot Stats
            self.__calculate_stat()
            self._print_last_stats()
            if (epoch + 1) % self.plot_step == 0:
                clear_output()
                self._plot_fake_images()

    def __train_disc(self, real_image, fake_images):
        """
        Epoch training discriminator
        Args:
            real_image: torch.Tensor (batch_size, image_channles, H, W)
            fake_image: the same
        """
        # Pass real images through discriminator
        real_pred = self.model["discriminator"](real_image)
        real_target = torch.ones(real_image.size(0), 1, 1, 1, device=self.device)
        real_loss = self.criterion["discriminator"](real_pred, real_target)
        self.score["real"][-1].append(torch.mean(real_pred.detach().cpu()))
        fake_pred = self.model["discriminator"](fake_images)
        fake_target = torch.zeros(fake_images.size(0), 1, 1, 1, device=self.device)
        fake_loss = self.criterion["discriminator"](fake_pred, fake_target)
        self.score["fake"][-1].append(torch.mean(fake_pred.detach().cpu()))
        loss_d = real_loss + fake_loss
        self.losses["discriminator"][-1].append(loss_d.detach().item())
        return loss_d

    def __train_gen(self, fake_images):
        """
        Epoch training generator
        Args:
            real_image: torch.Tensor (batch_size, image_channles, H, W)
            fake_image: the same
        """
        fake_image = self._get_fake_images(self.batch_size)
        pred = self.model["discriminator"](fake_image)
        target = torch.ones(self.batch_size, 1, 1, 1, device=self.device)
        loss_g = self.criterion["generator"](pred, target)
        self.losses["generator"][-1].append(loss_g.detach().item())
        return loss_g

    def __prepare_stat(self):
        """
        Aappend empty array to losses and scores history:
            [1, 2, 3, ..., [stats_per_epochs]]
        and at the end we calculate np.mean to get ONE number instead ARRAY
            [1, 2, 3, ..., []] -> [1, 2, 3, ..., n]
        """
        self.losses["discriminator"].append([])
        self.losses["generator"].append([])
        self.score["real"].append([])
        self.score["fake"].append([])

    def __calculate_stat(self):
        """
        [1, 2, 3, ..., [stats_per_epochs]] -> [1, 2, 3, ..., n]
        """
        self.losses["discriminator"][-1] = np.mean(self.losses["discriminator"][-1])
        self.losses["generator"][-1] = np.mean(self.losses["generator"][-1])
        self.score["real"][-1] = np.mean(self.score["real"][-1])
        self.score["fake"][-1] = np.mean(self.score["fake"][-1])

    def _get_fake_images(self, N):
        latent = torch.randn(N, self.latent_size, 1, 1, device=self.device)
        fake_images = self.model["generator"](latent)
        return fake_images

    def _print_last_stats(self):
        """
        Warning: you should use self.__calculate_stat() before this function
        """
        print(
            "Epoch {}, loss_g {:.4f}, loss_d {:.4f}, fake_score {:.4f}, real_score {:.4f}".format(
                len(self.score["real"]),
                self.losses["generator"][-1],
                self.losses["discriminator"][-1],
                self.score["fake"][-1],
                self.score["real"][-1],
            )
        )

    def _plot_fake_images(self, nrow=8):
        """
        Showing the generator's results
        """
        grid = torchvision.utils.make_grid(
            self._get_fake_images(nrow**2), nrow=nrow, padding=0, scale_each=True
        )
        fig = plt.figure(figsize=(16, 8))
        plt.imshow(grid.cpu().permute(1, 2, 0))
        plt.axis("off")
        plt.show()

    def _plot_stats(self):
        """
        Plotting stats of history training
        """
        fig, axes = plt.subplots(1, 2, figsize=(10, 4))
        sns.lineplot(self.losses["discriminator"], label="discriminator", ax=axes[0])
        sns.lineplot(self.losses["generator"], label="generator", ax=axes[0])
        sns.lineplot(self.score["real"], label="real", ax=axes[1])
        sns.lineplot(self.score["fake"], label="fake", ax=axes[1])
        plt.tight_layout()
        plt.show()


lr = 0.0001
models = {"discriminator": DCdisc, "generator": DCgen}
criterions = {"discriminator": nn.BCELoss(), "generator": nn.BCELoss()}
optimizers = {
    "discriminator": torch.optim.Adam(
        models["discriminator"].parameters(), lr=lr, betas=(0.5, 0.999)
    ),
    "generator": torch.optim.Adam(
        models["generator"].parameters(), lr=lr, betas=(0.5, 0.999)
    ),
}
trainer = GANTrainer(models, optimizers, criterions, latent_size, batch_size, device)
trainer._plot_fake_images()
trainer.train(faces_dataloader, 100)
trainer._plot_stats()


# Постройте графики лосса для генератора и дискриминатора. Что вы можете сказать про эти графики?
# ## Часть 3. Генерация изображений (1 балл)
# Теперь давайте оценим качество получившихся изображений. Напишите функцию, которая выводит изображения, сгенерированные нашим генератором
n_images = 4
fixed_latent = torch.randn(n_images, latent_size, 1, 1, device=device)
fake_images = model["generator"](fixed_latent)


def show_images(generated):
    # TODO: show generated images
    pass
<|endoftext|>
[{"Id": 132554970, "ScriptId": 39497739, "ParentScriptVersionId": NaN, "ScriptLanguageId": 9, "AuthorUserId": 3910696, "CreationDate": "06/06/2023 18:55:17", "VersionNumber": 3.0, "Title": "faceDCGAN", "EvaluationDate": "06/06/2023", "IsChange": true, "TotalLines": 599.0, "LinesInsertedFromPrevious": 86.0, "LinesChangedFromPrevious": 0.0, "LinesUnchangedFromPrevious": 513.0, "LinesInsertedFromFork": NaN, "LinesDeletedFromFork": NaN, "LinesChangedFromFork": NaN, "LinesUnchangedFromFork": NaN, "TotalVotes": 1}]
[{"Id": 190049945, "KernelVersionId": 132554970, "SourceDatasetVersionId": 3684316}]
[{"Id": 3684316, "DatasetId": 2204890, "DatasourceVersionId": 3738480, "CreatorUserId": 8975186, "LicenseName": "Attribution 4.0 International (CC BY 4.0)", "CreationDate": "05/23/2022 17:12:59", "VersionNumber": 1.0, "Title": "Flickr-Faces-HQ (FFHQ) small", "Slug": "faces-dataset-small", "Subtitle": "3143 photos of Flickr-Faces-HQ (1024px)", "Description": "This is a small version of FFHQ with 3143 photos. \nFlickr-Faces-HQ (FFHQ) is an image dataset containing high-quality images of human faces. It is provided by NVIDIA under the Creative Commons BY-NC-SA 4.0 license. It offers 70,000 PNG images at 1024\u00d71024 resolution that display diverse ages, ethnicities, image backgrounds, and accessories like hats and eyeglasses.", "VersionNotes": "Initial release", "TotalCompressedBytes": 0.0, "TotalUncompressedBytes": 0.0}]
[{"Id": 2204890, "CreatorUserId": 8975186, "OwnerUserId": 8975186.0, "OwnerOrganizationId": NaN, "CurrentDatasetVersionId": 3684316.0, "CurrentDatasourceVersionId": 3738480.0, "ForumId": 2231003, "Type": 2, "CreationDate": "05/23/2022 17:12:59", "LastActivityDate": "05/23/2022", "TotalViews": 2531, "TotalDownloads": 1206, "TotalVotes": 13, "TotalKernels": 12}]
#
# Школа глубокого обучения ФПМИ МФТИ
# Домашнее задание. Generative adversarial networks
# В этом домашнем задании вы обучите GAN генерировать лица людей и посмотрите на то, как можно оценивать качество генерации
import os
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
import torchvision.transforms as tt
import torchvision
import torch
import torch.nn as nn
import cv2
from tqdm.notebook import tqdm
from torchvision.utils import save_image
from torchvision.utils import make_grid
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from torch.utils.data import Dataset
import multiprocessing as mp
from PIL import Image
from IPython.display import clear_output

sns.set(style="darkgrid", font_scale=1.2)
sns.set_style("darkgrid")


# ## Часть 1. Подготовка данных (1 балл)
# В качестве обучающей выборки возьмем часть датасета [Flickr Faces](https://github.com/NVlabs/ffhq-dataset), который содержит изображения лиц людей в высоком разрешении (1024х1024). Оригинальный датасет очень большой, поэтому мы возьмем его часть. Скачать датасет можно [здесь](https://www.kaggle.com/datasets/tommykamaz/faces-dataset-small?resource=download-directory) и  [здесь](https://drive.google.com/drive/folders/14H7LQWzd09SaM11oLfGtd455sWH2rmRW?usp=share_link)
# Давайте загрузим наши изображения. Напишите функцию, которая строит DataLoader для изображений, при этом меняя их размер до нужного значения (размер 1024 слишком большой, поэтому мы рекомендуем взять размер 128 либо немного больше)
class FacesDataset(Dataset):
    """
    Dataset for Flickr-Faces-HQ (FFHQ) small
    """

    def __init__(self, path, input_transform=None):
        super().__init__()
        self.path = path
        self.input_transform = input_transform
        self.paths = os.listdir(self.path)
        self.data = []
        self._read_images()

    def _read_image(self, image_name):
        """
        Read image and use input_transforms
        """
        image = Image.open(os.path.join(self.path, image_name))
        if self.input_transform:
            image = self.input_transform(image)
        return image

    def _read_images(self):
        """
        Pool of _read_image functions
        """
        with mp.Pool(processes=mp.cpu_count()) as pool:
            self.data = list(
                tqdm(
                    pool.imap_unordered(self._read_image, self.paths),
                    total=len(self.paths),
                )
            )

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]


def to_device(data, device):
    """
    Move data to self.device
    """
    if isinstance(data, (list, tuple)):
        return [to_device(x, device) for x in data]
    return data.to(device, non_blocking=True)


class DeviceDataLoader:
    def __init__(self, dataloader, device):
        self.dataloader = dataloader
        self.device = device

    def __iter__(self):
        for batch in self.dataloader:
            yield to_device(batch, self.device)

    def __len__(self):
        return len(self.dataloader)


def get_dataloader(path, device, input_transform, image_size, batch_size):
    """
    Builds dataloader for training data.
    Use tt.Compose and tt.Resize for transformations
    :param image_size: height and wdith of the image
    :param batch_size: batch_size of the dataloader
    :returns: DataLoader object
    """
    # TODO: resize images, convert them to tensors and build dataloader
    faces_dataset = FacesDataset(path, input_transform)
    return DeviceDataLoader(
        DataLoader(faces_dataset, batch_size=batch_size, drop_last=True), device
    )


# !ls /kaggle/input/faces-dataset-small/faces_dataset_small
image_size = (128, 128)
batch_size = 64
path = "/kaggle/input/faces-dataset-small/faces_dataset_small"
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
input_transform = tt.Compose(
    [
        tt.RandomHorizontalFlip(p=0.5),
        tt.Resize(image_size),
        tt.ToTensor(),
        tt.ConvertImageDtype(torch.float),
    ]
)
faces_dataloader = get_dataloader(path, device, input_transform, image_size, batch_size)
faces = next(iter(faces_dataloader))
grid = torchvision.utils.make_grid(faces, nrow=8, padding=0, scale_each=True)
fig = plt.figure(figsize=(16, 8))
plt.imshow(grid.cpu().permute(1, 2, 0))
plt.axis("off")
plt.show()


# ## Часть 2. Построение и обучение модели (2 балла)
# Сконструируйте генератор и дискриминатор. Помните, что:
# * дискриминатор принимает на вход изображение (тензор размера `3 x image_size x image_size`) и выдает вероятность того, что изображение настоящее (тензор размера 1)
# * генератор принимает на вход тензор шумов размера `latent_size x 1 x 1` и генерирует изображение размера `3 x image_size x image_size`
# ### PixelNorm
# to avoid training BatchNorm module
class PixelNorm(nn.Module):
    """
    Nontrainable pixel normalization layer
    """

    def __init__(self, eps=1e-8):
        super(PixelNorm, self).__init__()
        self.eps = eps

    def forward(self, x):
        return (
            x
            / torch.sqrt(torch.mean(x**2, dim=1, keepdim=True) + self.eps)
            / x.shape[0]
        )


# Что же такое PixelNorm?
# В [статье по ProGAN]("https://towardsdatascience.com/progan-how-nvidia-generated-images-of-unprecedented-quality-51c98ec2cbd2") предлагается использовать такую нормализация вместо BatchNorm. Главным отличием от BatchNorm является отсутствие тренировочных весов у этого слоя.
# $$b_{x,y}=\frac{a_{x,y}}{\sqrt{\frac{1}{C}\sum\limits_{j=0}^C (a_{x,y}^j)^2 + \varepsilon}}$$
# Но как же влияет PixelNorm на фотографию?
faces = next(iter(faces_dataloader))
grid = torchvision.utils.make_grid(faces, nrow=8, padding=0, scale_each=True)
fig = plt.figure(figsize=(16, 8))
# Using PixelNorm
plt.imshow(PixelNorm()(grid).cpu().permute(1, 2, 0))
plt.axis("off")
plt.show()


# Норм, однако давайте не использовать нормализацию на последнем слое генератора, чтобы он смог восстанавливать цвета.
# ### Generator
class DCGenerator(nn.Module):
    """
    Generator of Deep Convolution GAN
    """

    def __init__(self, z_dim=96, channels_image=3, features_g=8):
        super(DCGenerator, self).__init__()
        self.channels_image = channels_image
        self.latent_size = z_dim
        self._block1 = self._block(
            z_dim, features_g * 32, 4, 1, 0
        )  # N x f_g*16 x 4 x 4
        self._block2 = self._block(
            features_g * 32, features_g * 16, 3, 2, 1, output_padding=1
        )  # 8 x 8
        self._block3 = self._block(
            features_g * 16, features_g * 8, 3, 2, 1, output_padding=1
        )  # 16 x 16
        self._block4 = self._block(
            features_g * 8, features_g * 4, 3, 2, 1, output_padding=1
        )  # 32 x 32
        self._block5 = self._block(features_g * 4, features_g * 2, 4, 2, 1)  # 64 x 64
        self.end_block = nn.Sequential(
            nn.ConvTranspose2d(
                features_g * 2, self.channels_image, 4, 2, 1, output_padding=0
            ),  # 128 x 128
            nn.Tanh(),  # [0, +inf] -> [0, 1]
        )

    def _block(
        self, in_channels, out_channels, kernel_size, stride, padding, output_padding=0
    ):
        """
        Generator block consist such elems as
            nn.ConvTranspose2d
            PixelNorm()
            nn.ReLU()
        """
        return nn.Sequential(
            nn.ConvTranspose2d(
                in_channels,
                out_channels,
                kernel_size,
                stride,
                padding,
                output_padding=output_padding,
                bias=False,
            ),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
        )

    def _get_input_tensor_shape(self):
        """
        Return input tensor size
        """
        return self.latent_size

    def forward(self, x):
        x = self._block1(x)
        x = self._block2(x)
        x = tt.GaussianBlur(3)(x)  # Blur Image
        x = self._block3(x)
        # x = tt.GaussianBlur(3)(x) # Blur image
        x = self._block4(x)
        x = self._block5(x)
        x = self.end_block(x)
        return x


# ### Disciminator
class DCDiscriminator(nn.Module):
    """
    Discriminator of Deep Convolution GAN
    """

    def __init__(self, channels_image, features_d, H, W):
        super(DCDiscriminator, self).__init__()
        self.channels_image = channels_image
        self.H = H
        self.W = W
        self.disc = nn.Sequential(
            # N x channels_image x 128 x 128
            nn.Conv2d(
                channels_image, features_d, kernel_size=4, stride=2, padding=1
            ),  # 64 x 64
            nn.LeakyReLU(0.2),
            self._block(features_d, features_d * 2, 4, 2, 1),  # 32 x 32
            self._block(features_d * 2, features_d * 4, 4, 2, 1),  # 16 x 16
            self._block(features_d * 4, features_d * 8, 4, 2, 1),  # 8 x 8
            self._block(features_d * 8, features_d * 16, 4, 2, 1),  # 4 x 4
            nn.Conv2d(features_d * 16, 1, kernel_size=4, stride=2, padding=0),  # 1 x 1
            nn.Sigmoid(),
        )

    def _block(self, in_channels, out_channels, kernel_size, stride, padding):
        """
        Discriminator block consist such elems as
            nn.Conv2d
            nn.BatchNorm2d
            nn.LeakyReLU
        """
        return nn.Sequential(
            nn.Conv2d(
                in_channels, out_channels, kernel_size, stride, padding, bias=False
            ),
            nn.BatchNorm2d(out_channels),
            nn.LeakyReLU(0.1),
        )

    def _get_input_tensor_shape(self):
        """
        Return input tensor size
        """
        return self.channels_image, self.H, self.W

    def forward(self, x):
        return self.disc(x)


# Проинизицализируем веса нормальным распределением с стандартным отлонением 0 и дисперсией 0.02
def initialize_weights(model):
    """
    Model  normal initialization
    Args:
        model: nn.Module object
    """
    for m in model.modules():
        if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.BatchNorm2d)):
            nn.init.normal_(m.weight, 0.0, 0.02)


# Проверим размерности выходов генератора и дискриминатора
def checkout(gen, disc, latent_size, N, in_channels, H, W):
    """
    Checking right shape of generator and discriminator
    Args:
        gen: initilized generator
        disc: initilized discriminator
        latent_size: lenght of latent vector
        N: batch size
        in_channels: source image channels
        (H, W): source image size
    """
    x = torch.randn((N, in_channels, H, W))
    assert disc(x).shape == (N, 1, 1, 1), "Discriminator has wrong output shape"
    z = torch.randn((N, latent_size, 1, 1))
    assert gen(z).shape == (
        N,
        in_channels,
        H,
        W,
    ), "Generator has wrong output shape: {}. \nExpected ({}, {}, {}, {})".format(
        gen(z).shape, N, in_channels, H, W
    )
    print("Success")


latent_size = 96  # choose latent size
batch_size = 8
image_channels = 3
H, W = 128, 128
features = 32
DCdisc = DCDiscriminator(image_channels, features, H, W)
initialize_weights(DCdisc)
DCgen = DCGenerator(latent_size, image_channels, features)
initialize_weights(DCgen)
# do some test
checkout(DCgen, DCdisc, latent_size, batch_size, image_channels, H, W)


# ### Trainer
# Перейдем теперь к обучению нашего GANа. Алгоритм обучения следующий:
# 1. Учим дискриминатор:
# * берем реальные изображения и присваиваем им метку 1
# * генерируем изображения генератором и присваиваем им метку 0
# * обучаем классификатор на два класса
# 2. Учим генератор:
# * генерируем изображения генератором и присваиваем им метку 0
# * предсказываем дискриминаторором, реальное это изображение или нет
# В качестве функции потерь берем бинарную кросс-энтропию
class GANTrainer:
    def __init__(
        self,
        models,
        optimizers,
        criterions,
        latent_size,
        batch_size,
        device,
        plot_step=10,
    ):
        self.device = device
        # init and move to device
        self.model = {name: model.to(self.device) for name, model in models.items()}
        self.criterion = {
            name: criterion.to(self.device) for name, criterion in criterions.items()
        }
        self.optimizer = optimizers
        self.latent_size = latent_size
        self.batch_size = batch_size
        # Plot settings
        self.plot_step = plot_step
        # Losses
        self.losses = {"discriminator": [], "generator": []}
        self.score = {"real": [0], "fake": [0]}

    def train(self, dataloader, epochs):
        """
        Main train GAN function
        """
        self.model["discriminator"].train()
        self.model["generator"].train()
        torch.cuda.empty_cache()
        for epoch in range(epochs):
            self.__prepare_stat()
            print("Epoch {} out of {}".format(epoch + 1, epochs))
            for i, real_image in enumerate(tqdm(dataloader)):
                # Train discriminator
                if i < 5 or self.score["real"][-2] < 0.9:
                    self.optimizer["discriminator"].zero_grad()  # Clear grads
                    loss_d = self.__train_disc(
                        real_image, self._get_fake_images(self.batch_size)
                    )
                    loss_d.backward()  # Calculate gradients
                    self.optimizer["discriminator"].step()  # Doing optimazer step
                # Train Generator
                self.optimizer["generator"].zero_grad()
                loss_g = self.__train_gen(self._get_fake_images(self.batch_size))
                loss_g.backward()
                self.optimizer["generator"].step()
            # Plot Stats
            self.__calculate_stat()
            self._print_last_stats()
            if (epoch + 1) % self.plot_step == 0:
                clear_output()
                self._plot_fake_images()

    def __train_disc(self, real_image, fake_images):
        """
        Epoch training discriminator
        Args:
            real_image: torch.Tensor (batch_size, image_channles, H, W)
            fake_image: the same
        """
        # Pass real images through discriminator
        real_pred = self.model["discriminator"](real_image)
        real_target = torch.ones(real_image.size(0), 1, 1, 1, device=self.device)
        real_loss = self.criterion["discriminator"](real_pred, real_target)
        self.score["real"][-1].append(torch.mean(real_pred.detach().cpu()))
        fake_pred = self.model["discriminator"](fake_images)
        fake_target = torch.zeros(fake_images.size(0), 1, 1, 1, device=self.device)
        fake_loss = self.criterion["discriminator"](fake_pred, fake_target)
        self.score["fake"][-1].append(torch.mean(fake_pred.detach().cpu()))
        loss_d = real_loss + fake_loss
        self.losses["discriminator"][-1].append(loss_d.detach().item())
        return loss_d

    def __train_gen(self, fake_images):
        """
        Epoch training generator
        Args:
            real_image: torch.Tensor (batch_size, image_channles, H, W)
            fake_image: the same
        """
        fake_image = self._get_fake_images(self.batch_size)
        pred = self.model["discriminator"](fake_image)
        target = torch.ones(self.batch_size, 1, 1, 1, device=self.device)
        loss_g = self.criterion["generator"](pred, target)
        self.losses["generator"][-1].append(loss_g.detach().item())
        return loss_g

    def __prepare_stat(self):
        """
        Aappend empty array to losses and scores history:
            [1, 2, 3, ..., [stats_per_epochs]]
        and at the end we calculate np.mean to get ONE number instead ARRAY
            [1, 2, 3, ..., []] -> [1, 2, 3, ..., n]
        """
        self.losses["discriminator"].append([])
        self.losses["generator"].append([])
        self.score["real"].append([])
        self.score["fake"].append([])

    def __calculate_stat(self):
        """
        [1, 2, 3, ..., [stats_per_epochs]] -> [1, 2, 3, ..., n]
        """
        self.losses["discriminator"][-1] = np.mean(self.losses["discriminator"][-1])
        self.losses["generator"][-1] = np.mean(self.losses["generator"][-1])
        self.score["real"][-1] = np.mean(self.score["real"][-1])
        self.score["fake"][-1] = np.mean(self.score["fake"][-1])

    def _get_fake_images(self, N):
        latent = torch.randn(N, self.latent_size, 1, 1, device=self.device)
        fake_images = self.model["generator"](latent)
        return fake_images

    def _print_last_stats(self):
        """
        Warning: you should use self.__calculate_stat() before this function
        """
        print(
            "Epoch {}, loss_g {:.4f}, loss_d {:.4f}, fake_score {:.4f}, real_score {:.4f}".format(
                len(self.score["real"]),
                self.losses["generator"][-1],
                self.losses["discriminator"][-1],
                self.score["fake"][-1],
                self.score["real"][-1],
            )
        )

    def _plot_fake_images(self, nrow=8):
        """
        Showing the generator's results
        """
        grid = torchvision.utils.make_grid(
            self._get_fake_images(nrow**2), nrow=nrow, padding=0, scale_each=True
        )
        fig = plt.figure(figsize=(16, 8))
        plt.imshow(grid.cpu().permute(1, 2, 0))
        plt.axis("off")
        plt.show()

    def _plot_stats(self):
        """
        Plotting stats of history training
        """
        fig, axes = plt.subplots(1, 2, figsize=(10, 4))
        sns.lineplot(self.losses["discriminator"], label="discriminator", ax=axes[0])
        sns.lineplot(self.losses["generator"], label="generator", ax=axes[0])
        sns.lineplot(self.score["real"], label="real", ax=axes[1])
        sns.lineplot(self.score["fake"], label="fake", ax=axes[1])
        plt.tight_layout()
        plt.show()


lr = 0.0001
models = {"discriminator": DCdisc, "generator": DCgen}
criterions = {"discriminator": nn.BCELoss(), "generator": nn.BCELoss()}
optimizers = {
    "discriminator": torch.optim.Adam(
        models["discriminator"].parameters(), lr=lr, betas=(0.5, 0.999)
    ),
    "generator": torch.optim.Adam(
        models["generator"].parameters(), lr=lr, betas=(0.5, 0.999)
    ),
}
trainer = GANTrainer(models, optimizers, criterions, latent_size, batch_size, device)
trainer._plot_fake_images()
trainer.train(faces_dataloader, 100)
trainer._plot_stats()


# Постройте графики лосса для генератора и дискриминатора. Что вы можете сказать про эти графики?
# ## Часть 3. Генерация изображений (1 балл)
# Теперь давайте оценим качество получившихся изображений. Напишите функцию, которая выводит изображения, сгенерированные нашим генератором
n_images = 4
fixed_latent = torch.randn(n_images, latent_size, 1, 1, device=device)
fake_images = model["generator"](fixed_latent)


def show_images(generated):
    # TODO: show generated images
    pass
