StyleGan 2 ආදර්ශ පුහුණුව

31import math
32from pathlib import Path
33from typing import Iterator, Tuple
34
35import torch
36import torch.utils.data
37import torchvision
38from PIL import Image
39
40from labml import tracker, lab, monit, experiment
41from labml.configs import BaseConfigs
42from labml_helpers.device import DeviceConfigs
43from labml_helpers.train_valid import ModeState, hook_model_outputs
44from labml_nn.gan.stylegan import Discriminator, Generator, MappingNetwork, GradientPenalty, PathLengthPenalty
45from labml_nn.gan.wasserstein import DiscriminatorLoss, GeneratorLoss
46from labml_nn.utils import cycle_dataloader

49class Dataset(torch.utils.data.Dataset):

56    def __init__(self, path: str, image_size: int):

61        super().__init__()

64        self.paths = [p for p in Path(path).glob(f'**/*.jpg')]

67        self.transform = torchvision.transforms.Compose([

69            torchvision.transforms.Resize(image_size),

71            torchvision.transforms.ToTensor(),
72        ])

74    def __len__(self):

76        return len(self.paths)

78    def __getitem__(self, index):

80        path = self.paths[index]
81        img = Image.open(path)
82        return self.transform(img)

85class Configs(BaseConfigs):

93    device: torch.device = DeviceConfigs()

96    discriminator: Discriminator

98    generator: Generator

100    mapping_network: MappingNetwork

104    discriminator_loss: DiscriminatorLoss
105    generator_loss: GeneratorLoss

108    generator_optimizer: torch.optim.Adam
109    discriminator_optimizer: torch.optim.Adam
110    mapping_network_optimizer: torch.optim.Adam

113    gradient_penalty = GradientPenalty()

115    gradient_penalty_coefficient: float = 10.

118    path_length_penalty: PathLengthPenalty

121    loader: Iterator

124    batch_size: int = 32

126    d_latent: int = 512

128    image_size: int = 32

130    mapping_network_layers: int = 8

132    learning_rate: float = 1e-3

134    mapping_network_learning_rate: float = 1e-5

136    gradient_accumulate_steps: int = 1

138    adam_betas: Tuple[float, float] = (0.0, 0.99)

140    style_mixing_prob: float = 0.9

143    training_steps: int = 150_000

146    n_gen_blocks: int

154    lazy_gradient_penalty_interval: int = 4

156    lazy_path_penalty_interval: int = 32

158    lazy_path_penalty_after: int = 5_000

161    log_generated_interval: int = 500

163    save_checkpoint_interval: int = 2_000

166    mode: ModeState

168    log_layer_outputs: bool = False

175    dataset_path: str = str(lab.get_data_path() / 'stylegan2')

177    def init(self):

182        dataset = Dataset(self.dataset_path, self.image_size)

184        dataloader = torch.utils.data.DataLoader(dataset, batch_size=self.batch_size, num_workers=8,
185                                                 shuffle=True, drop_last=True, pin_memory=True)

187        self.loader = cycle_dataloader(dataloader)

190        log_resolution = int(math.log2(self.image_size))

193        self.discriminator = Discriminator(log_resolution).to(self.device)
194        self.generator = Generator(log_resolution, self.d_latent).to(self.device)

196        self.n_gen_blocks = self.generator.n_blocks

198        self.mapping_network = MappingNetwork(self.d_latent, self.mapping_network_layers).to(self.device)

200        self.path_length_penalty = PathLengthPenalty(0.99).to(self.device)

203        if self.log_layer_outputs:
204            hook_model_outputs(self.mode, self.discriminator, 'discriminator')
205            hook_model_outputs(self.mode, self.generator, 'generator')
206            hook_model_outputs(self.mode, self.mapping_network, 'mapping_network')

209        self.discriminator_loss = DiscriminatorLoss().to(self.device)
210        self.generator_loss = GeneratorLoss().to(self.device)

213        self.discriminator_optimizer = torch.optim.Adam(
214            self.discriminator.parameters(),
215            lr=self.learning_rate, betas=self.adam_betas
216        )
217        self.generator_optimizer = torch.optim.Adam(
218            self.generator.parameters(),
219            lr=self.learning_rate, betas=self.adam_betas
220        )
221        self.mapping_network_optimizer = torch.optim.Adam(
222            self.mapping_network.parameters(),
223            lr=self.mapping_network_learning_rate, betas=self.adam_betas
224        )

227        tracker.set_image("generated", True)

229    def get_w(self, batch_size: int):

243        if torch.rand(()).item() < self.style_mixing_prob:

245            cross_over_point = int(torch.rand(()).item() * self.n_gen_blocks)

247            z2 = torch.randn(batch_size, self.d_latent).to(self.device)
248            z1 = torch.randn(batch_size, self.d_latent).to(self.device)

250            w1 = self.mapping_network(z1)
251            w2 = self.mapping_network(z2)

253            w1 = w1[None, :, :].expand(cross_over_point, -1, -1)
254            w2 = w2[None, :, :].expand(self.n_gen_blocks - cross_over_point, -1, -1)
255            return torch.cat((w1, w2), dim=0)

257        else:

259            z = torch.randn(batch_size, self.d_latent).to(self.device)

261            w = self.mapping_network(z)

263            return w[None, :, :].expand(self.n_gen_blocks, -1, -1)

265    def get_noise(self, batch_size: int):

272        noise = []

274        resolution = 4

277        for i in range(self.n_gen_blocks):

279            if i == 0:
280                n1 = None

282            else:
283                n1 = torch.randn(batch_size, 1, resolution, resolution, device=self.device)

285            n2 = torch.randn(batch_size, 1, resolution, resolution, device=self.device)

288            noise.append((n1, n2))

291            resolution *= 2

294        return noise

296    def generate_images(self, batch_size: int):

304        w = self.get_w(batch_size)

306        noise = self.get_noise(batch_size)

309        images = self.generator(w, noise)

312        return images, w

314    def step(self, idx: int):

320        with monit.section('Discriminator'):

322            self.discriminator_optimizer.zero_grad()

325            for i in range(self.gradient_accumulate_steps):

327                with self.mode.update(is_log_activations=(idx + 1) % self.log_generated_interval == 0):

329                    generated_images, _ = self.generate_images(self.batch_size)

331                    fake_output = self.discriminator(generated_images.detach())

334                    real_images = next(self.loader).to(self.device)

336                    if (idx + 1) % self.lazy_gradient_penalty_interval == 0:
337                        real_images.requires_grad_()

339                    real_output = self.discriminator(real_images)

342                    real_loss, fake_loss = self.discriminator_loss(real_output, fake_output)
343                    disc_loss = real_loss + fake_loss

346                    if (idx + 1) % self.lazy_gradient_penalty_interval == 0:

348                        gp = self.gradient_penalty(real_images, real_output)
349                        tracker.add('loss.gp', gp)

351                        disc_loss = disc_loss + 0.5 * self.gradient_penalty_coefficient * gp * self.lazy_gradient_penalty_interval

354                    disc_loss.backward()

357                    tracker.add('loss.discriminator', disc_loss)
358
359            if (idx + 1) % self.log_generated_interval == 0:

361                tracker.add('discriminator', self.discriminator)

364            torch.nn.utils.clip_grad_norm_(self.discriminator.parameters(), max_norm=1.0)

366            self.discriminator_optimizer.step()

369        with monit.section('Generator'):

371            self.generator_optimizer.zero_grad()
372            self.mapping_network_optimizer.zero_grad()

375            for i in range(self.gradient_accumulate_steps):

377                generated_images, w = self.generate_images(self.batch_size)

379                fake_output = self.discriminator(generated_images)

382                gen_loss = self.generator_loss(fake_output)

385                if idx > self.lazy_path_penalty_after and (idx + 1) % self.lazy_path_penalty_interval == 0:

387                    plp = self.path_length_penalty(w, generated_images)

389                    if not torch.isnan(plp):
390                        tracker.add('loss.plp', plp)
391                        gen_loss = gen_loss + plp

394                gen_loss.backward()

397                tracker.add('loss.generator', gen_loss)
398
399            if (idx + 1) % self.log_generated_interval == 0:

401                tracker.add('generator', self.generator)
402                tracker.add('mapping_network', self.mapping_network)

405            torch.nn.utils.clip_grad_norm_(self.generator.parameters(), max_norm=1.0)
406            torch.nn.utils.clip_grad_norm_(self.mapping_network.parameters(), max_norm=1.0)

409            self.generator_optimizer.step()
410            self.mapping_network_optimizer.step()

413        if (idx + 1) % self.log_generated_interval == 0:
414            tracker.add('generated', torch.cat([generated_images[:6], real_images[:3]], dim=0))

416        if (idx + 1) % self.save_checkpoint_interval == 0:
417            experiment.save_checkpoint()

420        tracker.save()

422    def train(self):

428        for i in monit.loop(self.training_steps):

430            self.step(i)

432            if (i + 1) % self.log_generated_interval == 0:
433                tracker.new_line()

436def main():

442    experiment.create(name='stylegan2')

444    configs = Configs()

447    experiment.configs(configs, {
448        'device.cuda_device': 0,
449        'image_size': 64,
450        'log_generated_interval': 200
451    })

454    configs.init()

456    experiment.add_pytorch_models(mapping_network=configs.mapping_network,
457                                  generator=configs.generator,
458                                  discriminator=configs.discriminator)

461    with experiment.start():

463        configs.train()

467if __name__ == '__main__':
468    main()