ai-art

 1import torch
 2import torch.nn as nn
 3
 4
 5class ActNorm(nn.Module):
 6    def __init__(self, num_features, logdet=False, affine=True, allow_reverse_init=False):
 7        assert affine
 8        super().__init__()
 9        self.logdet = logdet
10        self.loc = nn.Parameter(torch.zeros(1, num_features, 1, 1))
11        self.scale = nn.Parameter(torch.ones(1, num_features, 1, 1))
12        self.allow_reverse_init = allow_reverse_init
13
14        self.register_buffer('initialized', torch.tensor(0, dtype=torch.uint8))
15
16    def initialize(self, input):
17        with torch.no_grad():
18            flatten = input.permute(1, 0, 2, 3).contiguous().view(input.shape[1], -1)
19            mean = (
20                flatten.mean(1)
21                .unsqueeze(1)
22                .unsqueeze(2)
23                .unsqueeze(3)
24                .permute(1, 0, 2, 3)
25            )
26            std = (
27                flatten.std(1)
28                .unsqueeze(1)
29                .unsqueeze(2)
30                .unsqueeze(3)
31                .permute(1, 0, 2, 3)
32            )
33
34            self.loc.data.copy_(-mean)
35            self.scale.data.copy_(1 / (std + 1e-6))
36
37    def forward(self, input, reverse=False):
38        if reverse:
39            return self.reverse(input)
40        if len(input.shape) == 2:
41            input = input[:,:,None,None]
42            squeeze = True
43        else:
44            squeeze = False
45
46        _, _, height, width = input.shape
47
48        if self.training and self.initialized.item() == 0:
49            self.initialize(input)
50            self.initialized.fill_(1)
51
52        h = self.scale * (input + self.loc)
53
54        if squeeze:
55            h = h.squeeze(-1).squeeze(-1)
56
57        if self.logdet:
58            log_abs = torch.log(torch.abs(self.scale))
59            logdet = height*width*torch.sum(log_abs)
60            logdet = logdet * torch.ones(input.shape[0]).to(input)
61            return h, logdet
62
63        return h
64
65    def reverse(self, output):
66        if self.training and self.initialized.item() == 0:
67            if not self.allow_reverse_init:
68                raise RuntimeError(
69                    "Initializing ActNorm in reverse direction is "
70                    "disabled by default. Use allow_reverse_init=True to enable."
71                )
72            else:
73                self.initialize(output)
74                self.initialized.fill_(1)
75
76        if len(output.shape) == 2:
77            output = output[:,:,None,None]
78            squeeze = True
79        else:
80            squeeze = False
81
82        h = output / self.scale - self.loc
83
84        if squeeze:
85            h = h.squeeze(-1).squeeze(-1)
86        return h