<jupyter_start><jupyter_text>Model<jupyter_code>class Generator(nn.Module):
    def __init__(self):
        super().__init__()

        self.deconv1 = nn.Sequential(
            nn.ConvTranspose2d(128, 512, kernel_size=4, stride=1, padding=0),
            nn.BatchNorm2d(512),
            nn.ReLU())

        self.deconv2 = nn.Sequential(
            nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU())

        self.deconv3 = nn.Sequential(
            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU())

        self.deconv4 = nn.Sequential(
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU())

        self.deconv5 = nn.Sequential(
            nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1),
            nn.Tanh()
        )
    
    def forward(self, z):
        out = self.deconv1(z)
        out = self.deconv2(out)
        out = self.deconv3(out)
        out = self.deconv4(out)
        out = self.deconv5(out)
        return out
z = torch.randn(64, 128, 1, 1)
generator = Generator()
out = generator(z)
out.shape
class Discriminator(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU(0.2))

        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2))

        self.conv3 = nn.Sequential(
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2))

        self.conv4 = nn.Sequential(
            nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.2))
        
        self.conv5 = nn.Conv2d(512, 1, kernel_size=4, stride=1, padding=0)

    def forward(self, img):
        out = self.conv1(img)
        out = self.conv2(out)
        out = self.conv3(out)
        out = self.conv4(out)
        logit = self.conv5(out)
        return logit
discriminator = Discriminator()
logit = discriminator(out)
logit.shape<jupyter_output><empty_output><jupyter_text>Loss<jupyter_code>def d_loss_fn(r_logit, f_logit):
    return -r_logit.mean() + f_logit.mean()

def g_loss_fn(f_logit):
    return - f_logit.mean()<jupyter_output><empty_output><jupyter_text>Data<jupyter_code>from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import CelebA
transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.CenterCrop(148),
    transforms.Resize(64),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
train_dataset = CelebA(root='../../data', split='train', transform=transform, download=False)
val_dataset = CelebA(root='../../data', split='test', transform=transform, download=False)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, drop_last=True)
val_loader = DataLoader(val_dataset, batch_size=144, shuffle=True, drop_last=True)
batch = next(iter(train_loader))
batch[0].shape, batch[1].shape<jupyter_output><empty_output><jupyter_text>Training<jupyter_code>import torch.optim as optim
G_optimizer = torch.optim.Adam(generator.parameters(), lr=0.0002)
D_optimizer = torch.optim.Adam(discriminator.parameters(), lr=0.0002)<jupyter_output><empty_output><jupyter_text>Training discriminator<jupyter_code>x_real = next(iter(train_loader))[0]
x_real.shape
# Training discriminator
z = torch.randn(32, 128, 1, 1)
print(z.shape)
x_fake = generator(z).detach()
print(x_fake.shape)
x_real_d_logit = discriminator(x_real)
x_real_d_logit.shape
x_fake_d_logit = discriminator(x_fake)
x_fake_d_logit.shape
d_loss = d_loss_fn(x_real_d_logit, x_fake_d_logit)
d_loss<jupyter_output><empty_output><jupyter_text>Gradient Penalty<jupyter_code>x_real.shape, x_fake.shape
# realとfakeの補間画像を返す
def sample(real, fake):
    shape = [real.shape[0], 1, 1, 1]
    alpha = torch.rand(shape, device=real.device)
    sample = alpha * fake + (1 - alpha) * real
    return sample
x_sample = sample(x_real, x_fake)
x_sample.shape
x_sample
import functools

# TODO: このpartialはなぜ必要？ないとgradでエラーが出る
pred = functools.partial(discriminator)(x_sample)
pred.shape
x_sample.requires_grad = True
grad = torch.autograd.grad(pred, x_sample, grad_outputs=torch.ones_like(pred), create_graph=True)[0]
grad.shape
grad_norm = grad.view(grad.size(0), -1).norm(p=2, dim=1)
grad_norm.shape
gradient_penalty = ((grad_norm - 1)** 2).mean()
gradient_penalty<jupyter_output><empty_output><jupyter_text>Training Generator<jupyter_code>z = torch.randn(32, 128, 1, 1)
x_fake = generator(z)
x_fake.shape
x_fake_d_logit = discriminator(x_fake)
x_fake_d_logit.shape
g_loss = g_loss_fn(x_fake_d_logit)
g_loss<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13351877607788595, 'fraction_numerical': 0.04232068348897278, 'mean_word_length': 3.63963133640553, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '371845', 'license': 'MIT', 'n_tokens_mistral': 2016, 'n_tokens_neox': 1884, 'n_words': 295, 'path': 'notebooks/wgan_gp_celeba/model.ipynb', 'repo_name': 'aidiary/vae-pytorch', 'size': '5033'}