Styleගන්2

148import math
149from typing import Tuple, Optional, List
150
151import numpy as np
152import torch
153import torch.nn.functional as F
154import torch.utils.data
155from torch import nn

158class MappingNetwork(nn.Module):

173    def __init__(self, features: int, n_layers: int):

178        super().__init__()

181        layers = []
182        for i in range(n_layers):

184            layers.append(EqualizedLinear(features, features))

186            layers.append(nn.LeakyReLU(negative_slope=0.2, inplace=True))
187
188        self.net = nn.Sequential(*layers)

190    def forward(self, z: torch.Tensor):

192        z = F.normalize(z, dim=1)

194        return self.net(z)

197class Generator(nn.Module):

214    def __init__(self, log_resolution: int, d_latent: int, n_features: int = 32, max_features: int = 512):

221        super().__init__()

226        features = [min(max_features, n_features * (2 ** i)) for i in range(log_resolution - 2, -1, -1)]

228        self.n_blocks = len(features)

231        self.initial_constant = nn.Parameter(torch.randn((1, features[0], 4, 4)))

234        self.style_block = StyleBlock(d_latent, features[0], features[0])
235        self.to_rgb = ToRGB(d_latent, features[0])

238        blocks = [GeneratorBlock(d_latent, features[i - 1], features[i]) for i in range(1, self.n_blocks)]
239        self.blocks = nn.ModuleList(blocks)

243        self.up_sample = UpSample()

245    def forward(self, w: torch.Tensor, input_noise: List[Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]]):

255        batch_size = w.shape[1]

258        x = self.initial_constant.expand(batch_size, -1, -1, -1)

261        x = self.style_block(x, w[0], input_noise[0][1])

263        rgb = self.to_rgb(x, w[0])

266        for i in range(1, self.n_blocks):

268            x = self.up_sample(x)

270            x, rgb_new = self.blocks[i - 1](x, w[i], input_noise[i])

272            rgb = self.up_sample(rgb) + rgb_new

275        return rgb

278class GeneratorBlock(nn.Module):

294    def __init__(self, d_latent: int, in_features: int, out_features: int):

300        super().__init__()

303        self.style_block1 = StyleBlock(d_latent, in_features, out_features)

305        self.style_block2 = StyleBlock(d_latent, out_features, out_features)

308        self.to_rgb = ToRGB(d_latent, out_features)

310    def forward(self, x: torch.Tensor, w: torch.Tensor, noise: Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]):

318        x = self.style_block1(x, w, noise[0])

321        x = self.style_block2(x, w, noise[1])

324        rgb = self.to_rgb(x, w)

327        return x, rgb

330class StyleBlock(nn.Module):

344    def __init__(self, d_latent: int, in_features: int, out_features: int):

350        super().__init__()

353        self.to_style = EqualizedLinear(d_latent, in_features, bias=1.0)

355        self.conv = Conv2dWeightModulate(in_features, out_features, kernel_size=3)

357        self.scale_noise = nn.Parameter(torch.zeros(1))

359        self.bias = nn.Parameter(torch.zeros(out_features))

362        self.activation = nn.LeakyReLU(0.2, True)

364    def forward(self, x: torch.Tensor, w: torch.Tensor, noise: Optional[torch.Tensor]):

371        s = self.to_style(w)

373        x = self.conv(x, s)

375        if noise is not None:
376            x = x + self.scale_noise[None, :, None, None] * noise

378        return self.activation(x + self.bias[None, :, None, None])

381class ToRGB(nn.Module):

394    def __init__(self, d_latent: int, features: int):

399        super().__init__()

402        self.to_style = EqualizedLinear(d_latent, features, bias=1.0)

405        self.conv = Conv2dWeightModulate(features, 3, kernel_size=1, demodulate=False)

407        self.bias = nn.Parameter(torch.zeros(3))

409        self.activation = nn.LeakyReLU(0.2, True)

411    def forward(self, x: torch.Tensor, w: torch.Tensor):

417        style = self.to_style(w)

419        x = self.conv(x, style)

421        return self.activation(x + self.bias[None, :, None, None])

424class Conv2dWeightModulate(nn.Module):

431    def __init__(self, in_features: int, out_features: int, kernel_size: int,
432                 demodulate: float = True, eps: float = 1e-8):

440        super().__init__()

442        self.out_features = out_features

444        self.demodulate = demodulate

446        self.padding = (kernel_size - 1) // 2

449        self.weight = EqualizedWeight([out_features, in_features, kernel_size, kernel_size])

451        self.eps = eps

453    def forward(self, x: torch.Tensor, s: torch.Tensor):

460        b, _, h, w = x.shape

463        s = s[:, None, :, None, None]

465        weights = self.weight()[None, :, :, :, :]

470        weights = weights * s

473        if self.demodulate:

475            sigma_inv = torch.rsqrt((weights ** 2).sum(dim=(2, 3, 4), keepdim=True) + self.eps)

477            weights = weights * sigma_inv

480        x = x.reshape(1, -1, h, w)

483        _, _, *ws = weights.shape
484        weights = weights.reshape(b * self.out_features, *ws)

488        x = F.conv2d(x, weights, padding=self.padding, groups=b)

491        return x.reshape(-1, self.out_features, h, w)

494class Discriminator(nn.Module):

508    def __init__(self, log_resolution: int, n_features: int = 64, max_features: int = 512):

514        super().__init__()

517        self.from_rgb = nn.Sequential(
518            EqualizedConv2d(3, n_features, 1),
519            nn.LeakyReLU(0.2, True),
520        )

525        features = [min(max_features, n_features * (2 ** i)) for i in range(log_resolution - 1)]

527        n_blocks = len(features) - 1

529        blocks = [DiscriminatorBlock(features[i], features[i + 1]) for i in range(n_blocks)]
530        self.blocks = nn.Sequential(*blocks)

533        self.std_dev = MiniBatchStdDev()

535        final_features = features[-1] + 1

537        self.conv = EqualizedConv2d(final_features, final_features, 3)

539        self.final = EqualizedLinear(2 * 2 * final_features, 1)

541    def forward(self, x: torch.Tensor):

547        x = x - 0.5

549        x = self.from_rgb(x)

551        x = self.blocks(x)

554        x = self.std_dev(x)

556        x = self.conv(x)

558        x = x.reshape(x.shape[0], -1)

560        return self.final(x)

563class DiscriminatorBlock(nn.Module):

574    def __init__(self, in_features, out_features):

579        super().__init__()

581        self.residual = nn.Sequential(DownSample(),
582                                      EqualizedConv2d(in_features, out_features, kernel_size=1))

585        self.block = nn.Sequential(
586            EqualizedConv2d(in_features, in_features, kernel_size=3, padding=1),
587            nn.LeakyReLU(0.2, True),
588            EqualizedConv2d(in_features, out_features, kernel_size=3, padding=1),
589            nn.LeakyReLU(0.2, True),
590        )

593        self.down_sample = DownSample()

596        self.scale = 1 / math.sqrt(2)

598    def forward(self, x):

600        residual = self.residual(x)

603        x = self.block(x)

605        x = self.down_sample(x)

608        return (x + residual) * self.scale

611class MiniBatchStdDev(nn.Module):

623    def __init__(self, group_size: int = 4):

627        super().__init__()
628        self.group_size = group_size

630    def forward(self, x: torch.Tensor):

635        assert x.shape[0] % self.group_size == 0

638        grouped = x.view(self.group_size, -1)

645        std = torch.sqrt(grouped.var(dim=0) + 1e-8)

647        std = std.mean().view(1, 1, 1, 1)

649        b, _, h, w = x.shape
650        std = std.expand(b, -1, h, w)

652        return torch.cat([x, std], dim=1)

655class DownSample(nn.Module):

667    def __init__(self):
668        super().__init__()

670        self.smooth = Smooth()

672    def forward(self, x: torch.Tensor):

674        x = self.smooth(x)

676        return F.interpolate(x, (x.shape[2] // 2, x.shape[3] // 2), mode='bilinear', align_corners=False)

679class UpSample(nn.Module):

690    def __init__(self):
691        super().__init__()

693        self.up_sample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)

695        self.smooth = Smooth()

697    def forward(self, x: torch.Tensor):

699        return self.smooth(self.up_sample(x))

702class Smooth(nn.Module):

711    def __init__(self):
712        super().__init__()

714        kernel = [[1, 2, 1],
715                  [2, 4, 2],
716                  [1, 2, 1]]

718        kernel = torch.tensor([[kernel]], dtype=torch.float)

720        kernel /= kernel.sum()

722        self.kernel = nn.Parameter(kernel, requires_grad=False)

724        self.pad = nn.ReplicationPad2d(1)

726    def forward(self, x: torch.Tensor):

728        b, c, h, w = x.shape

730        x = x.view(-1, 1, h, w)

733        x = self.pad(x)

736        x = F.conv2d(x, self.kernel)

739        return x.view(b, c, h, w)

742class EqualizedLinear(nn.Module):

751    def __init__(self, in_features: int, out_features: int, bias: float = 0.):

758        super().__init__()

760        self.weight = EqualizedWeight([out_features, in_features])

762        self.bias = nn.Parameter(torch.ones(out_features) * bias)

764    def forward(self, x: torch.Tensor):

766        return F.linear(x, self.weight(), bias=self.bias)

769class EqualizedConv2d(nn.Module):

778    def __init__(self, in_features: int, out_features: int,
779                 kernel_size: int, padding: int = 0):

786        super().__init__()

788        self.padding = padding

790        self.weight = EqualizedWeight([out_features, in_features, kernel_size, kernel_size])

792        self.bias = nn.Parameter(torch.ones(out_features))

794    def forward(self, x: torch.Tensor):

796        return F.conv2d(x, self.weight(), bias=self.bias, padding=self.padding)

799class EqualizedWeight(nn.Module):

820    def __init__(self, shape: List[int]):

824        super().__init__()

827        self.c = 1 / math.sqrt(np.prod(shape[1:]))

829        self.weight = nn.Parameter(torch.randn(shape))

832    def forward(self):

834        return self.weight * self.c

837class GradientPenalty(nn.Module):

853    def forward(self, x: torch.Tensor, d: torch.Tensor):

860        batch_size = x.shape[0]

866        gradients, *_ = torch.autograd.grad(outputs=d,
867                                            inputs=x,
868                                            grad_outputs=d.new_ones(d.shape),
869                                            create_graph=True)

872        gradients = gradients.reshape(batch_size, -1)

874        norm = gradients.norm(2, dim=-1)

876        return torch.mean(norm ** 2)

879class PathLengthPenalty(nn.Module):

903    def __init__(self, beta: float):

907        super().__init__()

910        self.beta = beta

912        self.steps = nn.Parameter(torch.tensor(0.), requires_grad=False)

916        self.exp_sum_a = nn.Parameter(torch.tensor(0.), requires_grad=False)

918    def forward(self, w: torch.Tensor, x: torch.Tensor):

925        device = x.device

927        image_size = x.shape[2] * x.shape[3]

929        y = torch.randn(x.shape, device=device)

933        output = (x * y).sum() / math.sqrt(image_size)

936        gradients, *_ = torch.autograd.grad(outputs=output,
937                                            inputs=w,
938                                            grad_outputs=torch.ones(output.shape, device=device),
939                                            create_graph=True)

942        norm = (gradients ** 2).sum(dim=2).mean(dim=1).sqrt()

945        if self.steps > 0:

948            a = self.exp_sum_a / (1 - self.beta ** self.steps)

952            loss = torch.mean((norm - a) ** 2)
953        else:

955            loss = norm.new_tensor(0)

958        mean = norm.mean().detach()

960        self.exp_sum_a.mul_(self.beta).add_(mean, alpha=1 - self.beta)

962        self.steps.add_(1.)

965        return loss