Skip to content

equalized_lr

Blur 

Blur(blur_kernel: List[int], factor: int, kernel_size: int)

Bases: nn.Module

Upsample (factor > 0)

Applied after a transpose convolution of stride U and kernel size K

Apply blurring FIR filter (before / after) a (downsampling / upsampling) op

Parameters:

Name Type Description Default
input Tensor

(N, C, (H - 1) * U + K - 1 + 1, (W - 1) * U + K - 1 + 1)

 required
blur_kernel Tensor

FIR filter

 required
factor int

U. Defaults to 2.

 required
kernel_size int

K. Defaults to 3.

 required

Returns:

Name Type Description
Tensor

(N, C, H * U, W * U)
Downsample (factor < 0)

Applied before a convolution of stride U and kernel size K

Parameters:

Name Type Description Default
input Tensor

(N, C, H, W)

 required
blur_kernel Tensor

FIR filter

 required
factor int

U. Defaults to 2.

 required
kernel_size int

K. Defaults to 3.

 required

Returns:

Name Type Description
Tensor

(N, C, H - (U + 1) + K - 1, H - (U + 1) + K - 1)
Source code in stylegan2_torch/equalized_lr.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
def __init__(self, blur_kernel: List[int], factor: int, kernel_size: int):
    """
    Apply blurring FIR filter (before / after) a (downsampling / upsampling) op

    Case 1: Upsample (factor > 0)
        Applied after a transpose convolution of stride U and kernel size K

    Args:
        input (Tensor): (N, C, (H - 1) * U + K - 1 + 1, (W - 1) * U + K - 1 + 1)
        blur_kernel (Tensor): FIR filter
        factor (int, optional): U. Defaults to 2.
        kernel_size (int, optional): K. Defaults to 3.

    Returns:
        Tensor: (N, C, H * U, W * U)


    Case 2: Downsample (factor < 0)
        Applied before a convolution of stride U and kernel size K

    Args:
        input (Tensor): (N, C, H, W)
        blur_kernel (Tensor): FIR filter
        factor (int, optional): U. Defaults to 2.
        kernel_size (int, optional): K. Defaults to 3.

    Returns:
        Tensor: (N, C, H - (U + 1) + K - 1, H  - (U + 1) + K - 1)
    """
    super().__init__()

    if factor > 0:
        p = (len(blur_kernel) - factor) - (kernel_size - 1)
        pad0 = (p + 1) // 2 + factor - 1
        pad1 = p // 2 + 1
    else:
        p = (len(blur_kernel) - abs(factor)) + (kernel_size - 1)
        pad0 = (p + 1) // 2
        pad1 = p // 2

    # Factor to compensate for averaging with zeros if upsampling
    self.kernel: Tensor
    self.register_buffer(
        "kernel", make_kernel(blur_kernel, factor if factor > 0 else 1)
    )
    self.pad = (pad0, pad1)

__call__ class-attribute 

__call__ = proxy(forward)

kernel instance-attribute 

kernel: Tensor = None

pad instance-attribute 

pad = (pad0, pad1)

forward 

forward(input: Tensor) -> Tensor
Source code in stylegan2_torch/equalized_lr.py 
171
172
def forward(self, input: Tensor) -> Tensor:
    return upfirdn2d(input, self.kernel, pad=self.pad)

EqualConv2d 

EqualConv2d(
    in_channel: int,
    out_channel: int,
    kernel_size: int,
    stride: int = 1,
    padding: int = 0,
    bias: bool = True,
)

Bases: nn.Module

Conv2d with equalized learning rate

Source code in stylegan2_torch/equalized_lr.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def __init__(
    self,
    in_channel: int,
    out_channel: int,
    kernel_size: int,
    stride: int = 1,
    padding: int = 0,
    bias: bool = True,
):
    super().__init__()

    # Equalized Learning Rate
    self.weight = Parameter(
        torch.randn(out_channel, in_channel, kernel_size, kernel_size)
    )
    # std = gain / sqrt(fan_in)
    self.scale = 1 / math.sqrt(in_channel * kernel_size**2)
    self.stride = stride
    self.padding = padding
    self.bias = Parameter(torch.zeros(out_channel)) if bias else None

__call__ class-attribute 

__call__ = proxy(forward)

bias instance-attribute 

bias = Parameter(torch.zeros(out_channel)) if bias else None

padding instance-attribute 

padding = padding

scale instance-attribute 

scale = 1 / math.sqrt(in_channel * kernel_size ** 2)

stride instance-attribute 

stride = stride

weight instance-attribute 

weight = Parameter(
    torch.randn(
        out_channel, in_channel, kernel_size, kernel_size
    )
)

__repr__ 

__repr__() -> str
Source code in stylegan2_torch/equalized_lr.py 
51
52
53
54
55
def __repr__(self) -> str:
    return (
        f"{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},"
        f" {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})"
    )

forward 

forward(input: Tensor) -> Tensor
Source code in stylegan2_torch/equalized_lr.py 
42
43
44
45
46
47
48
49
def forward(self, input: Tensor) -> Tensor:
    return conv2d(
        input=input,
        weight=self.weight * self.scale,
        bias=self.bias,
        stride=self.stride,
        padding=self.padding,
    )

EqualLeakyReLU 

EqualLeakyReLU(
    in_dim: int, out_dim: int, lr_mult: float = 1
)

Bases: nn.Module

Leaky ReLU with equalized learning rate

Source code in stylegan2_torch/equalized_lr.py 
 99
100
101
102
103
104
105
106
107
108
109
def __init__(self, in_dim: int, out_dim: int, lr_mult: float = 1):
    super().__init__()

    # Equalized Learning Rate
    self.weight = Parameter(torch.randn(out_dim, in_dim).div_(lr_mult))

    self.bias = Parameter(torch.zeros(out_dim))

    self.scale = (1 / math.sqrt(in_dim)) * lr_mult

    self.lr_mult = lr_mult

__call__ class-attribute 

__call__ = proxy(forward)

bias instance-attribute 

bias = Parameter(torch.zeros(out_dim))

lr_mult instance-attribute 

lr_mult = lr_mult

scale instance-attribute 

scale = 1 / math.sqrt(in_dim) * lr_mult

weight instance-attribute 

weight = Parameter(
    torch.randn(out_dim, in_dim).div_(lr_mult)
)

__repr__ 

__repr__() -> str
Source code in stylegan2_torch/equalized_lr.py 
115
116
117
118
def __repr__(self) -> str:
    return (
        f"{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})"
    )

forward 

forward(input: Tensor) -> Tensor
Source code in stylegan2_torch/equalized_lr.py 
111
112
113
def forward(self, input: Tensor) -> Tensor:
    out = F.linear(input, self.weight * self.scale)
    return fused_leaky_relu(out, self.bias * self.lr_mult)

EqualLinear 

EqualLinear(
    in_dim: int,
    out_dim: int,
    bias_init: int = 0,
    lr_mult: float = 1,
)

Bases: nn.Module

Linear with equalized learning rate

Source code in stylegan2_torch/equalized_lr.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def __init__(
    self,
    in_dim: int,
    out_dim: int,
    bias_init: int = 0,
    lr_mult: float = 1,
):
    super().__init__()

    # Equalized Learning Rate
    self.weight = Parameter(torch.randn(out_dim, in_dim).div_(lr_mult))

    self.bias = Parameter(torch.zeros(out_dim).fill_(bias_init))

    self.scale = (1 / math.sqrt(in_dim)) * lr_mult

    self.lr_mult = lr_mult

__call__ class-attribute 

__call__ = proxy(forward)

bias instance-attribute 

bias = Parameter(torch.zeros(out_dim).fill_(bias_init))

lr_mult instance-attribute 

lr_mult = lr_mult

scale instance-attribute 

scale = 1 / math.sqrt(in_dim) * lr_mult

weight instance-attribute 

weight = Parameter(
    torch.randn(out_dim, in_dim).div_(lr_mult)
)

__repr__ 

__repr__() -> str
Source code in stylegan2_torch/equalized_lr.py 
86
87
88
89
def __repr__(self) -> str:
    return (
        f"{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})"
    )

forward 

forward(input: Tensor) -> Tensor
Source code in stylegan2_torch/equalized_lr.py 
83
84
def forward(self, input: Tensor) -> Tensor:
    return F.linear(input, self.weight * self.scale, bias=self.bias * self.lr_mult)