ස්ථාවර විසරණය සඳහා ස්වයංක්රීය ආකේතකය

18from typing import List
19
20import torch
21import torch.nn.functional as F
22from torch import nn

25class Autoencoder(nn.Module):

32    def __init__(self, encoder: 'Encoder', decoder: 'Decoder', emb_channels: int, z_channels: int):

39        super().__init__()
40        self.encoder = encoder
41        self.decoder = decoder

44        self.quant_conv = nn.Conv2d(2 * z_channels, 2 * emb_channels, 1)

47        self.post_quant_conv = nn.Conv2d(emb_channels, z_channels, 1)

49    def encode(self, img: torch.Tensor) -> 'GaussianDistribution':

56        z = self.encoder(img)

58        moments = self.quant_conv(z)

60        return GaussianDistribution(moments)

62    def decode(self, z: torch.Tensor):

69        z = self.post_quant_conv(z)

71        return self.decoder(z)

74class Encoder(nn.Module):

79    def __init__(self, *, channels: int, channel_multipliers: List[int], n_resnet_blocks: int,
80                 in_channels: int, z_channels: int):

89        super().__init__()

93        n_resolutions = len(channel_multipliers)

96        self.conv_in = nn.Conv2d(in_channels, channels, 3, stride=1, padding=1)

99        channels_list = [m * channels for m in [1] + channel_multipliers]

102        self.down = nn.ModuleList()

104        for i in range(n_resolutions):

106            resnet_blocks = nn.ModuleList()

108            for _ in range(n_resnet_blocks):
109                resnet_blocks.append(ResnetBlock(channels, channels_list[i + 1]))
110                channels = channels_list[i + 1]

112            down = nn.Module()
113            down.block = resnet_blocks

115            if i != n_resolutions - 1:
116                down.downsample = DownSample(channels)
117            else:
118                down.downsample = nn.Identity()

120            self.down.append(down)

123        self.mid = nn.Module()
124        self.mid.block_1 = ResnetBlock(channels, channels)
125        self.mid.attn_1 = AttnBlock(channels)
126        self.mid.block_2 = ResnetBlock(channels, channels)

129        self.norm_out = normalization(channels)
130        self.conv_out = nn.Conv2d(channels, 2 * z_channels, 3, stride=1, padding=1)

132    def forward(self, img: torch.Tensor):

138        x = self.conv_in(img)

141        for down in self.down:

143            for block in down.block:
144                x = block(x)

146            x = down.downsample(x)

149        x = self.mid.block_1(x)
150        x = self.mid.attn_1(x)
151        x = self.mid.block_2(x)

154        x = self.norm_out(x)
155        x = swish(x)
156        x = self.conv_out(x)

159        return x

162class Decoder(nn.Module):

167    def __init__(self, *, channels: int, channel_multipliers: List[int], n_resnet_blocks: int,
168                 out_channels: int, z_channels: int):

177        super().__init__()

181        num_resolutions = len(channel_multipliers)

184        channels_list = [m * channels for m in channel_multipliers]

187        channels = channels_list[-1]

190        self.conv_in = nn.Conv2d(z_channels, channels, 3, stride=1, padding=1)

193        self.mid = nn.Module()
194        self.mid.block_1 = ResnetBlock(channels, channels)
195        self.mid.attn_1 = AttnBlock(channels)
196        self.mid.block_2 = ResnetBlock(channels, channels)

199        self.up = nn.ModuleList()

201        for i in reversed(range(num_resolutions)):

203            resnet_blocks = nn.ModuleList()

205            for _ in range(n_resnet_blocks + 1):
206                resnet_blocks.append(ResnetBlock(channels, channels_list[i]))
207                channels = channels_list[i]

209            up = nn.Module()
210            up.block = resnet_blocks

212            if i != 0:
213                up.upsample = UpSample(channels)
214            else:
215                up.upsample = nn.Identity()

217            self.up.insert(0, up)

220        self.norm_out = normalization(channels)
221        self.conv_out = nn.Conv2d(channels, out_channels, 3, stride=1, padding=1)

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

229        h = self.conv_in(z)

232        h = self.mid.block_1(h)
233        h = self.mid.attn_1(h)
234        h = self.mid.block_2(h)

237        for up in reversed(self.up):

239            for block in up.block:
240                h = block(h)

242            h = up.upsample(h)

245        h = self.norm_out(h)
246        h = swish(h)
247        img = self.conv_out(h)

250        return img

253class GaussianDistribution:

258    def __init__(self, parameters: torch.Tensor):

264        self.mean, log_var = torch.chunk(parameters, 2, dim=1)

266        self.log_var = torch.clamp(log_var, -30.0, 20.0)

268        self.std = torch.exp(0.5 * self.log_var)

270    def sample(self):

272        return self.mean + self.std * torch.randn_like(self.std)

275class AttnBlock(nn.Module):

280    def __init__(self, channels: int):

284        super().__init__()

286        self.norm = normalization(channels)

288        self.q = nn.Conv2d(channels, channels, 1)
289        self.k = nn.Conv2d(channels, channels, 1)
290        self.v = nn.Conv2d(channels, channels, 1)

292        self.proj_out = nn.Conv2d(channels, channels, 1)

294        self.scale = channels ** -0.5

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

301        x_norm = self.norm(x)

303        q = self.q(x_norm)
304        k = self.k(x_norm)
305        v = self.v(x_norm)

310        b, c, h, w = q.shape
311        q = q.view(b, c, h * w)
312        k = k.view(b, c, h * w)
313        v = v.view(b, c, h * w)

316        attn = torch.einsum('bci,bcj->bij', q, k) * self.scale
317        attn = F.softmax(attn, dim=2)

320        out = torch.einsum('bij,bcj->bci', attn, v)

323        out = out.view(b, c, h, w)

325        out = self.proj_out(out)

328        return x + out

331class UpSample(nn.Module):

335    def __init__(self, channels: int):

339        super().__init__()

341        self.conv = nn.Conv2d(channels, channels, 3, padding=1)

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

348        x = F.interpolate(x, scale_factor=2.0, mode="nearest")

350        return self.conv(x)

353class DownSample(nn.Module):

357    def __init__(self, channels: int):

361        super().__init__()

363        self.conv = nn.Conv2d(channels, channels, 3, stride=2, padding=0)

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

370        x = F.pad(x, (0, 1, 0, 1), mode="constant", value=0)

372        return self.conv(x)

375class ResnetBlock(nn.Module):

379    def __init__(self, in_channels: int, out_channels: int):

384        super().__init__()

386        self.norm1 = normalization(in_channels)
387        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, stride=1, padding=1)

389        self.norm2 = normalization(out_channels)
390        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, stride=1, padding=1)

392        if in_channels != out_channels:
393            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, 1, stride=1, padding=0)
394        else:
395            self.nin_shortcut = nn.Identity()

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

402        h = x

405        h = self.norm1(h)
406        h = swish(h)
407        h = self.conv1(h)

410        h = self.norm2(h)
411        h = swish(h)
412        h = self.conv2(h)

415        return self.nin_shortcut(x) + h

418def swish(x: torch.Tensor):

424    return x * torch.sigmoid(x)

427def normalization(channels: int):

433    return nn.GroupNorm(num_groups=32, num_channels=channels, eps=1e-6)