<jupyter_start><jupyter_text>Progressive GAN<jupyter_code>import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import Adam
from torchvision.datasets import CIFAR10
from torchvision.transforms import Compose, ToTensor, Resize, Lambda
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import numpy as np<jupyter_output><empty_output><jupyter_text>Data<jupyter_code>def transform_in_size(size):
    return Compose([Resize((size, size)), ToTensor(), Lambda(lambda x: x * 2 - 1)])

def get_dataloader(size):
    dataset = CIFAR10('~/pytorch', download=True, transform=transform_in_size(size))
    return DataLoader(dataset, batch_size=32, drop_last=True, shuffle=True)<jupyter_output><empty_output><jupyter_text>CUDA<jupyter_code>if torch.cuda.is_available():
    device = 'cuda'
else:
    device = 'cpu'
    
print(f"Learning on: {device}")<jupyter_output><empty_output><jupyter_text>Generator* TODO: Change UpsamplingBilinear to interpolate()* TODO: Move to inner class<jupyter_code>class GeneratorBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.upsample = nn.UpsamplingBilinear2d(scale_factor=2)
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)
        
    def forward(self, x):
        _, x = self.forward_with_img(x)
        
        return x
    
    def forward_with_img(self, x):
        img = self.upsample(x)
        x = F.leaky_relu(self.conv1(img))
        x = self.bn1(x)
        x = F.leaky_relu(self.conv2(x))
        x = self.bn2(x)
        
        return img, x
class Generator(nn.Module):
    def __init__(self, channels=512):
        super().__init__()
        self.channels = channels
        self.trained_blocks = nn.ModuleList()
        self.new_block = self.create_initial_layer(channels)
        self.rgb = self.to_rgb(channels)
        
    def to_rgb(self, in_channels):
        return nn.Sequential(
            nn.Conv2d(in_channels, 3, 1, padding=0),
            nn.Tanh())
    
    def create_initial_layer(self, channels):
        module = nn.Sequential(
            nn.Conv2d(channels, channels, 4, padding=3),
            nn.BatchNorm2d(channels),
            nn.LeakyReLU(),
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.BatchNorm2d(channels),
            nn.LeakyReLU())
        
        return module
    
    def append_layer(self, channels):
        self.trained_blocks.append(self.new_block)
        self.new_block = GeneratorBlock(self.channels, channels)
        self.rgb = self.to_rgb(channels)
        self.channels = channels
        
    def forward(self, x, alpha=1.0):
        for block in self.trained_blocks:
            x = block(x)
            
        if alpha < 1.0:
            img , x = self.new_block.forward_with_img(x)
            img = self.rgb(img)
            x = self.rgb(x)
            
            x = img * (1 - alpha) + img * alpha
        else:
            x = self.new_block(x)
            x = self.rgb(x)
        
        return x<jupyter_output><empty_output><jupyter_text>Disciminator<jupyter_code>class DiscriminatorBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.downsample = nn.AvgPool2d(2, 2)
        
    def forward(self, x):
        x = F.leaky_relu(self.conv1(x))
        x = self.bn1(x)
        x = F.leaky_relu(self.conv2(x))
        x = self.bn2(x)
        x = self.downsample(x)
        
        return x
class Discriminator(nn.Module):
    def __init__(self, channels=512):
        super().__init__()
        self.in_channels = channels
        self.out_channels = channels
        self.trained_layers = nn.ModuleList()
        self.new_layer = self.create_initial_layer(channels)
        self.downsample = nn.AvgPool2d(2, 2)
        self.rgb = self.from_rgb(channels)
        self.rgb_skip = self.from_rgb(channels)
        
    def create_initial_layer(self, channels):
        module = nn.Sequential(
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.BatchNorm2d(channels),
            nn.LeakyReLU(),
            nn.Conv2d(channels, channels, 4, padding=0),
            nn.LeakyReLU(),
            nn.Flatten(),
            nn.Linear(channels, 1))
        
        return module
    
    def prepend_layer(self, channels):
        self.out_channels = self.in_channels
        self.in_channels = channels
        
        self.trained_layers.insert(0, self.new_layer)
        self.new_layer = DiscriminatorBlock(self.in_channels, self.out_channels)
        self.rgb = self.from_rgb(self.in_channels)
        self.rgb_skip = self.from_rgb(self.out_channels)
        
    def from_rgb(self, channels):
        return nn.Conv2d(3, channels, 1, padding=0)
        
    def forward(self, x, alpha=1.0):
        if alpha < 1.0:
            skip = self.downsample(x)
            skip = self.rgb_skip(skip)

            x = self.rgb(x)
            x = self.new_layer(x)
            x = skip * (1-alpha) + x * alpha
        else:
            x = self.rgb(x)
            x = self.new_layer(x)
        
        for layer in self.trained_layers:
            x = layer(x)
        
        return x
def train_batch():
    z = torch.randn(32, 512, 1, 1)
    gen_img = g(z)
    print(f"Image size: {gen_img.shape}")
    out = d(gen_img)
    
    print(out.shape)
def wgan_loss(y, target):
    return - (y * target).mean()
def clip_parameters(parameters):
    for p in parameters:
        p.data.clamp_(-0.01, 0.01)
def train_epoch(g, d, g_optim, d_optim, size, alpha=1.0):
    d_criterion = nn.BCEWithLogitsLoss()

    
    dl = get_dataloader(size)
    real = torch.ones((32, 1)).to(device)
    fake = -torch.ones((32, 1)).to(device)
    g_loss_total = 0
    d_loss_total = 0
    it = 0
    for real_img, y in dl:
        # print(real_img.shape)
        real_img = real_img.to(device)
        out = d(real_img, alpha)
        real_loss = wgan_loss(out, real) # d_criterion(out, real)
        fake_img = g(torch.randn((32, 512, 1, 1), device=device))
        out = d(fake_img, alpha)
        fake_loss = wgan_loss(out, fake) # d_criterion(out, fake)
        d_loss = 0.5 * (real_loss + fake_loss)
        
        d_optim.zero_grad()
        d_loss.backward()
        d_optim.step()
        clip_parameters(d.parameters())
        
        
        fake_img = g(torch.randn((32, 512, 1, 1), device=device))
        out = d(fake_img, alpha)
        g_loss = wgan_loss(out, real) # d_criterion(out, real)
        
        g_optim.zero_grad()
        g_loss.backward()
        g_optim.step()
        clip_parameters(g.parameters())
        
        it += 1
        g_loss_total += g_loss.item()
        d_loss_total += d_loss.item()
        if it % 400 == 0:
            print(f"a: {alpha:.2f} G_LOSS: {g_loss_total / it:.4f} - D_LOSS: {d_loss_total / it:.4f}")
            g_loss_total = 0
            d_loss_total = 0
            it = 0
            
        
def train_new_layer(g, d, size, g_layer, d_layer):
    g.append_layer(g_layer)
    d.prepend_layer(d_layer)
    g = g.to(device)
    d = d.to(device)
    
    d_optim = Adam(d.parameters(), lr=0.001, betas=(0.0, 0.99))
    g_optim = Adam(g.parameters(), lr=0.001, betas=(0.0, 0.99))
    
    for alpha in np.linspace(0.0, 1.0, 25):
        for i in range(1):
            train_epoch(g, d, g_optim, d_optim, size, alpha)
            
def train_initial_layer(g, d):
    d_optim = Adam(d.parameters(), lr=0.001, betas=(0.0, 0.99))
    g_optim = Adam(g.parameters(), lr=0.001, betas=(0.0, 0.99))
    
    for i in range(10):
        train_epoch(g, d, g_optim, d_optim, 4)
g = Generator().to(device)
d = Discriminator().to(device)
train_initial_layer(g, d)
train_new_layer(g, d, 8, 512, 256)
train_new_layer(g, d, 16, 512, 128)
train_new_layer(g, d, 32, 512, 64)
def show_image(img):
    img = (img + 1) / 2
    img = img.squeeze()
    img = img.permute(1, 2, 0)
    image = img.cpu().detach().numpy()
    plt.imshow(image)
    
g = g.eval()
d = d.eval()
image = g(torch.randn((100, 512, 1, 1), device=device))
out = d(image)
f = out.sigmoid() > 0.5
print(f)
print(f.float().sum() / 100)
show_image(image[91])
dl = get_dataloader(32)
for i, _ in dl:
    image = i.to(device)
    show_image(i[1])
    break<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12637741046831955, 'fraction_numerical': 0.026515151515151516, 'mean_word_length': 2.3408742331288344, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 22, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '597732', 'license': 'MIT', 'n_tokens_mistral': 3236, 'n_tokens_neox': 3105, 'n_words': 604, 'path': 'PGAN - Progressive GAN.ipynb', 'repo_name': 'stante/gan-architectures-pytorch', 'size': '8712'}