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

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

34class Configs(BaseConfigs):

41    device: torch.device = DeviceConfigs()

44    eps_model: UNet

46    diffusion: DenoiseDiffusion

49    image_channels: int = 3

51    image_size: int = 32

53    n_channels: int = 64

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

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

61    n_steps: int = 1_000

63    batch_size: int = 64

65    n_samples: int = 16

67    learning_rate: float = 2e-5

70    epochs: int = 1_000

73    dataset: torch.utils.data.Dataset

75    data_loader: torch.utils.data.DataLoader

78    optimizer: torch.optim.Adam

80    def init(self):

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

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

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

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

102        tracker.set_image("sample", True)

104    def sample(self):

108        with torch.no_grad():

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

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

116                t = self.n_steps - t_ - 1

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

121            tracker.save('sample', x)

123    def train(self):

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

131            tracker.add_global_step()

133            data = data.to(self.device)

136            self.optimizer.zero_grad()

138            loss = self.diffusion.loss(data)

140            loss.backward()

142            self.optimizer.step()

144            tracker.save('loss', loss)

146    def run(self):

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

152            self.train()

154            self.sample()

156            tracker.new_line()

158            experiment.save_checkpoint()

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

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

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

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

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

180    def __len__(self):

184        return len(self._files)

186    def __getitem__(self, index: int):

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

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

199    return CelebADataset(c.image_size)

202class MNISTDataset(torchvision.datasets.MNIST):

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

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

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

224    return MNISTDataset(c.image_size)

227def main():

229    experiment.create(name='diffuse', writers={'screen', 'labml'})

232    configs = Configs()

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

242    configs.init()

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

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

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