GAN_Artem

import numpy as np
import os
from PIL import Image
import matplotlib.pyplot as plt
from IPython.display import clear_output

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

from torchvision import transforms

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

! wget http://vis-www.cs.umass.edu/lfw/lfw-deepfunneled.tgz -O tmp.tgz
os.system("tar xvzf tmp.tgz && rm tmp.tgz")

images = []

for dirpath, _, filenames in os.walk('lfw-deepfunneled'):
    for fname in filenames:
        if fname.endswith(".jpg"):
            fpath = os.path.join(dirpath,fname)
            image = Image.open(fpath)
            images.append(image)
len(images)


def plot_gallery(images, n_row=3, n_col=6, from_torch=False):
    """Helper function to plot a gallery of portraits"""

    if from_torch:
        images = [x.data.numpy().transpose(1, 2, 0) for x in images]

    plt.figure(figsize=(1.5 * n_col, 1.7 * n_row))
    plt.subplots_adjust(bottom=0, left=.01, right=.99, top=.90, hspace=.35)

    for i in range(n_row * n_col):
        plt.subplot(n_row, n_col, i + 1)
        plt.imshow(images[i])

        plt.xticks(())
        plt.yticks(())

    plt.show()
plot_gallery(images)


class Faces(Dataset):
    def __init__(self, faces):
        self.data = faces
        self.transform = transforms.Compose([
                                    transforms.CenterCrop((90, 90)),
                                    transforms.Resize((64, 64)),
                                    transforms.ToTensor(),
                                    # transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))
                                    ])

    def __getitem__(self, index):

        x = self.data[index]
        return self.transform(x).float()

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

dataset = Faces(images)

# dataset[0] — это вызов метода __getitem__(0)
img = dataset[0]

print(img.shape)
plot_gallery(dataset, from_torch=True)


train_size = int(len(dataset) * 0.8)
val_size = len(dataset) - train_size

g_cpu = torch.Generator().manual_seed(8888)
train_data, val_data = torch.utils.data.random_split(dataset, [train_size, val_size], generator=g_cpu)
train_loader = torch.utils.data.DataLoader(train_data, batch_size=16, shuffle=True)
val_loader = torch.utils.data.DataLoader(val_data, batch_size=16, shuffle=False)

# Реализация GAN
# Дискриминатор и генератор - свёрточные сети.
class Discriminator(nn.Module):
    def __init__(self):
        super().__init__()
        # in: 3 x 64 x 64

        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        # out: 64 x 32 x 32

        self.conv2 = nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(128)
        # out: 128 x 16 x 16

        self.conv3 = nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1, bias=False)
        self.bn3 = nn.BatchNorm2d(256)
        # out: 256 x 8 x 8

        self.conv4 = nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1, bias=False)
        self.bn4 = nn.BatchNorm2d(512)
        # out: 512 x 4 x 4

        self.conv5 = nn.Conv2d(512, 1, kernel_size=4, stride=1, padding=0, bias=False)
        # out: 1 x 1 x 1

        self.flatten = nn.Flatten()
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):

        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = F.relu(self.bn4(self.conv4(x)))

        x = self.conv5(x)
        x = self.flatten(x)
        x = self.sigmoid(x)

        return x

class Generator(nn.Module):
    def __init__(self, latent_size):
        super().__init__()
        # in: latent_size x 1 x 1

        self.convT1 = nn.ConvTranspose2d(latent_size, 512, kernel_size=4, stride=1, padding=0, bias=False)
        self.bn1 = nn.BatchNorm2d(512)
        # out: 512 x 4 x 4

        self.convT2 = nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(256)
        # out: 256 x 8 x 8

        self.convT3 = nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=False)
        self.bn3 = nn.BatchNorm2d(128)
        # out: 128 x 16 x 16

        self.convT4 = nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, bias=False)
        self.bn4 = nn.BatchNorm2d(64)
        # out: 64 x 32 x 32

        self.convT5 = nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1, bias=False)
        self.tanh = nn.Tanh()
        # out: 3 x 64 x 64

    def forward(self, x):

        x = F.relu(self.bn1(self.convT1(x)))
        x = F.relu(self.bn2(self.convT2(x)))
        x = F.relu(self.bn3(self.convT3(x)))
        x = F.relu(self.bn4(self.convT4(x)))

        x = self.convT5(x)
        x = self.tanh(x)

        return x

# размер случайных векторов, подаваемых генератору для генерации изображений
latent_size = 128

# Заведем фиксированный вектор шума, чтобы в процессе обучения сети отслеживать то, как меняется вывод генератора на этом фиксированном # векторе:
fixed_latent = torch.randn(64, latent_size, 1, 1, device=device)

Обучение GAN
Перейдём теперь к тому, как обучается GAN. Процесс устроен как чередование двух шагов — обучения дискриминатора и генератора.

1. Обновляем дискриминатор.
Берём батч реальных изображений и ставим им метку 1.
Генерируем батч фейковых изображений с помощью генератора и назначаем им метку 0.
Тренируем дискриминатор как обычный бинарный классификатор, чтобы он максимально хорошо различал эти два типа примеров.

2. Обновляем генератор.
Замораживаем веса дискриминатора.
Генерируем изображения из случайного шума и пропускаем их через дискриминатор.
В функции потерь считаем, что для этих изображений целевая метка = 1 (мы хотим, чтобы дискриминатор принял их за реальные) и по этому сигналу обновляем только генератор.

from IPython.display import clear_output
from torchvision.utils import make_grid

# функция для де-нормализации
stats = (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)
def denorm(img_tensors):
    return img_tensors * stats[1][0] + stats[0][0]

# функция для отрисовки результатов генерации генератором
# на фиксированном векторе шума
def show_samples(latent_tensors):
    fake_images = generator(latent_tensors)

    fig, ax = plt.subplots(figsize=(8, 8))
    ax.set_xticks([]); ax.set_yticks([])
    ax.imshow(make_grid(denorm(fake_images[:64].cpu().detach()), nrow=8).permute(1, 2, 0))
    plt.show()

# функция непосредственного обучения
def train(models, opts, loss_fns, epochs, train_loader, val_loader, batch_size=64):

    # Losses & scores
    losses_g = []
    losses_d = []
    real_scores = []
    fake_scores = []

    for epoch in range(epochs):
        print('* Epoch %d/%d' % (epoch+1, epochs))

        # Обучаем сеть на изображениях из train_loader
        models["discriminator"].train()
        models["generator"].train()

        # массивы для сохранения значений лоссов и метрик accuracy
        # дискриминатора в процессе обучения сети
        loss_d_per_epoch = []
        loss_g_per_epoch = []
        real_score_per_epoch = []
        fake_score_per_epoch = []

        for i, X_batch in enumerate(train_loader):
            X_batch = X_batch.to(device)

            # 1. Шаг обучения дискриминатора.

            # Очистим градиенты дискриминатора
            opts["discriminator"].zero_grad()

            # 1.1 Получаем лосс на реальных данных
            # Подаем реальные изображения из батча X_batch на вход дискриминатору
            real_preds = models["discriminator"](X_batch)
            # Правильный ответ дискриминатора на эти изображения должен быть
            # вектором из единиц
            real_targets = torch.ones(X_batch.size(0), 1, device=device)
            # считаем значение лосса дискриминатора на батче реальных изображений
            real_loss = loss_fns["discriminator"](real_preds, real_targets)
            cur_real_score = torch.mean(real_preds).item()

            # 1.2 Получаем лосс на фейковых изображениях

            # Сгенерируем фейковые изображениях. Для этого сгенерируем случайный вектор шума
            # размера (batch_size, latent_size, 1, 1) и подадим на вход генератору.
            latent = torch.randn(batch_size, latent_size, 1, 1, device=device)
            # Получаем выход генератора на батч случайных векторов
            fake_images = models["generator"](latent)

            # Подаем фейковые изображениях из батча X_batch на вход дискриминатору
            fake_preds = models["discriminator"](fake_images)
            # Правильный ответ дискриминатора на эти изображениях должен быть
            # вектором из нулей
            fake_targets = torch.zeros(fake_images.size(0), 1, device=device)
            # считаем значение лосса дискриминатора на батче фейковых изображениях
            fake_loss = loss_fns["discriminator"](fake_preds, fake_targets)
            cur_fake_score = torch.mean(fake_preds).item()

            real_score_per_epoch.append(cur_real_score)
            fake_score_per_epoch.append(cur_fake_score)

            # 1.3 Обновляем веса дискриминатора: делаем шаг градиентного спуска
            loss_d = real_loss + fake_loss
            loss_d.backward()
            opts["discriminator"].step()
            loss_d_per_epoch.append(loss_d.item())

            # 2. Шаг обучения генератора
            # Очистим градиенты генератора
            opts["generator"].zero_grad()

            # Сгенерируем фейковые изображениях. Для этого сгенерируем случайный вектор шума
            # размера (batch_size, latent_size, 1, 1) и подадим на вход генератору.
            latent = torch.randn(batch_size, latent_size, 1, 1, device=device)
            # Получаем выход генератора на батч случайных векторов
            fake_images = models["generator"](latent)

            # Подаем фейковые изображениях из батча X_batch на вход дискриминатору
            preds = models["discriminator"](fake_images)
            # Задаем вектор "правильных" ответов на эти изображениях как вектор из единиц
            targets = torch.ones(batch_size, 1, device=device)
            # считаем лосс между preds и targets
            loss_g = loss_fns["generator"](preds, targets)

            # Обновляем веса генератора
            loss_g.backward()
            opts["generator"].step()
            loss_g_per_epoch.append(loss_g.item())

            # каждые 100 итераций обучения будем выводить текущие значения метрик
            # и отрисовывать результаты генерации генератором картинок
            # из фиксированных случайных векторов
            if i%100 == 0:
                # Record losses & scores
                losses_g.append(np.mean(loss_g_per_epoch))
                losses_d.append(np.mean(loss_d_per_epoch))
                real_scores.append(np.mean(real_score_per_epoch))
                fake_scores.append(np.mean(fake_score_per_epoch))

                # Log losses & scores (last batch)
                print("Epoch [{}/{}], loss_g: {:.4f}, loss_d: {:.4f}, real_score: {:.4f}, fake_score: {:.4f}".format(
                    epoch+1, epochs,
                    losses_g[-1], losses_d[-1], real_scores[-1], fake_scores[-1]))

                # Show generated images
                clear_output(wait=True)
                show_samples(fixed_latent)

    return losses_g, losses_d, real_scores, fake_scores


generator = Generator(latent_size).to(device)
discriminator = Discriminator().to(device)

models = {
    'generator': generator,
    'discriminator': discriminator
}

criterions = {
    "discriminator": nn.BCELoss(),
    "generator": nn.BCELoss()
}

lr = 0.0002
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))
    }

losses_g, losses_d, real_scores, fake_scores = train(models, optimizers, criterions, 10, train_loader, val_loader)

plt.figure(figsize=(15, 6))
plt.plot(losses_d, '-')
plt.plot(losses_g, '-')
plt.xlabel('epoch')
plt.ylabel('loss')
plt.legend(['Discriminator', 'Generator'])
plt.title('Losses');

plt.figure(figsize=(15, 6))

plt.plot(real_scores, '-')
plt.plot(fake_scores, '-')
plt.xlabel('epoch')
plt.ylabel('score')
plt.legend(['Real', 'Fake'])
plt.title('Scores');