විසරණ සම්භාවිතා ආකෘති (ඩීඩීපීඑම්) පුහුණුව නිරූපණය කිරීම

21from typing import List
22
23import torch
24import torch.utils.data
25import torchvision
26from PIL import Image
27
28from labml import lab, tracker, experiment, monit
29from labml.configs import BaseConfigs, option
30from labml_helpers.device import DeviceConfigs
31from labml_nn.diffusion.ddpm import DenoiseDiffusion
32from labml_nn.diffusion.ddpm.unet import UNet

35class Configs(BaseConfigs):

42    device: torch.device = DeviceConfigs()

45    eps_model: UNet

47    diffusion: DenoiseDiffusion

50    image_channels: int = 3

52    image_size: int = 32

54    n_channels: int = 64

57    channel_multipliers: List[int] = [1, 2, 2, 4]

59    is_attention: List[int] = [False, False, False, True]

62    n_steps: int = 1_000

64    batch_size: int = 64

66    n_samples: int = 16

68    learning_rate: float = 2e-5

71    epochs: int = 1_000

74    dataset: torch.utils.data.Dataset

76    data_loader: torch.utils.data.DataLoader

79    optimizer: torch.optim.Adam

81    def init(self):

83        self.eps_model = UNet(
84            image_channels=self.image_channels,
85            n_channels=self.n_channels,
86            ch_mults=self.channel_multipliers,
87            is_attn=self.is_attention,
88        ).to(self.device)

91        self.diffusion = DenoiseDiffusion(
92            eps_model=self.eps_model,
93            n_steps=self.n_steps,
94            device=self.device,
95        )

98        self.data_loader = torch.utils.data.DataLoader(self.dataset, self.batch_size, shuffle=True, pin_memory=True)

100        self.optimizer = torch.optim.Adam(self.eps_model.parameters(), lr=self.learning_rate)

103        tracker.set_image("sample", True)

105    def sample(self):

109        with torch.no_grad():

111            x = torch.randn([self.n_samples, self.image_channels, self.image_size, self.image_size],
112                            device=self.device)

115            for t_ in monit.iterate('Sample', self.n_steps):

117                t = self.n_steps - t_ - 1

119                x = self.diffusion.p_sample(x, x.new_full((self.n_samples,), t, dtype=torch.long))

122            tracker.save('sample', x)

124    def train(self):

130        for data in monit.iterate('Train', self.data_loader):

132            tracker.add_global_step()

134            data = data.to(self.device)

137            self.optimizer.zero_grad()

139            loss = self.diffusion.loss(data)

141            loss.backward()

143            self.optimizer.step()

145            tracker.save('loss', loss)

147    def run(self):

151        for _ in monit.loop(self.epochs):

153            self.train()

155            self.sample()

157            tracker.new_line()

159            experiment.save_checkpoint()

162class CelebADataset(torch.utils.data.Dataset):

167    def __init__(self, image_size: int):
168        super().__init__()

171        folder = lab.get_data_path() / 'celebA'

173        self._files = [p for p in folder.glob(f'**/*.jpg')]

176        self._transform = torchvision.transforms.Compose([
177            torchvision.transforms.Resize(image_size),
178            torchvision.transforms.ToTensor(),
179        ])

181    def __len__(self):

185        return len(self._files)

187    def __getitem__(self, index: int):

191        img = Image.open(self._files[index])
192        return self._transform(img)

195@option(Configs.dataset, 'CelebA')
196def celeb_dataset(c: Configs):

200    return CelebADataset(c.image_size)

203class MNISTDataset(torchvision.datasets.MNIST):

208    def __init__(self, image_size):
209        transform = torchvision.transforms.Compose([
210            torchvision.transforms.Resize(image_size),
211            torchvision.transforms.ToTensor(),
212        ])
213
214        super().__init__(str(lab.get_data_path()), train=True, download=True, transform=transform)

216    def __getitem__(self, item):
217        return super().__getitem__(item)[0]

220@option(Configs.dataset, 'MNIST')
221def mnist_dataset(c: Configs):

225    return MNISTDataset(c.image_size)

228def main():

230    experiment.create(name='diffuse', writers={'screen', 'comet'})

233    configs = Configs()

236    experiment.configs(configs, {
237        'dataset': 'CelebA',  # 'MNIST'
238        'image_channels': 3,  # 1,
239        'epochs': 100,  # 5,
240    })

243    configs.init()

246    experiment.add_pytorch_models({'eps_model': configs.eps_model})

249    with experiment.start():
250        configs.run()

254if __name__ == '__main__':
255    main()