StyleGAN 2

This is a PyTorch implementation of the paper Analyzing and Improving the Image Quality of StyleGAN which introduces StyleGAN 2. StyleGAN 2 is an improvement over StyleGAN from the paper A Style-Based Generator Architecture for Generative Adversarial Networks. And StyleGAN is based on Progressive GAN from the paper Progressive Growing of GANs for Improved Quality, Stability, and Variation. All three papers are from the same authors from NVIDIA AI.

Our implementation is a minimalistic StyleGAN 2 model training code. Only single GPU training is supported to keep the implementation simple. We managed to shrink it to keep it at less than 500 lines of code, including the training loop.

🏃 Here’s the training code: experiment.py.

These are $64 \times 64$ images generated after training for about 80K steps.

We’ll first introduce the three papers at a high level.

Generative Adversarial Networks

Generative adversarial networks have two components; the generator and the discriminator. The generator network takes a random latent vector ($z \in \mathcal{Z}$) and tries to generate a realistic image. The discriminator network tries to differentiate the real images from generated images. When we train the two networks together the generator starts generating images indistinguishable from real images.

Progressive GAN

Progressive GAN generates high-resolution images ($1080 \times 1080$) of size. It does so by progressively increasing the image size. First, it trains a network that produces a $4 \times 4$ image, then $8 \times 8$ , then an $16 \times 16$ image, and so on up to the desired image resolution.

At each resolution, the generator network produces an image in latent space which is converted into RGB, with a $1 \times 1$ convolution. When we progress from a lower resolution to a higher resolution (say from $4 \times 4$ to $8 \times 8$ ) we scale the latent image by $2\times$ and add a new block (two $3 \times 3$ convolution layers) and a new $1 \times 1$ layer to get RGB. The transition is done smoothly by adding a residual connection to the $2\times$ scaled $4 \times 4$ RGB image. The weight of this residual connection is slowly reduced, to let the new block take over.

The discriminator is a mirror image of the generator network. The progressive growth of the discriminator is done similarly.

$2\times$ and $0.5\times$ denote feature map resolution scaling and scaling. $4\times4$, $8\times4$, … denote feature map resolution at the generator or discriminator block. Each discriminator and generator block consists of 2 convolution layers with leaky ReLU activations.

They use minibatch standard deviation to increase variation and equalized learning rate which we discussed below in the implementation. They also use pixel-wise normalization where at each pixel the feature vector is normalized. They apply this to all the convolution layer outputs (except RGB).

StyleGAN

StyleGAN improves the generator of Progressive GAN keeping the discriminator architecture the same.

Mapping Network

It maps the random latent vector ($z \in \mathcal{Z}$) into a different latent space ($w \in \mathcal{W}$), with an 8-layer neural network. This gives an intermediate latent space $\mathcal{W}$ where the factors of variations are more linear (disentangled).

AdaIN

Then $w$ is transformed into two vectors (styles) per layer, $i$, $y_i = (y_{s,i}, y_{b,i}) = f_{A_i}(w)$ and used for scaling and shifting (biasing) in each layer with $\text{AdaIN}$ operator (normalize and scale): $$\text{AdaIN}(x_i, y_i) = y_{s, i} \frac{x_i - \mu(x_i)}{\sigma(x_i)} + y_{b,i}$$

Style Mixing

To prevent the generator from assuming adjacent styles are correlated, they randomly use different styles for different blocks. That is, they sample two latent vectors $(z_1, z_2)$ and corresponding $(w_1, w_2)$ and use $w_1$ based styles for some blocks and $w_2$ based styles for some blacks randomly.

Stochastic Variation

Noise is made available to each block which helps the generator create more realistic images. Noise is scaled per channel by a learned weight.

Bilinear Up and Down Sampling

All the up and down-sampling operations are accompanied by bilinear smoothing.

$A$ denotes a linear layer. $B$ denotes a broadcast and scaling operation (noise is a single channel). StyleGAN also uses progressive growing like Progressive GAN

StyleGAN 2

StyleGAN 2 changes both the generator and the discriminator of StyleGAN.

Weight Modulation and Demodulation

They remove the $\text{AdaIN}$ operator and replace it with the weight modulation and demodulation step. This is supposed to improve what they call droplet artifacts that are present in generated images, which are caused by the normalization in $\text{AdaIN}$ operator. Style vector per layer is calculated from $w_i \in \mathcal{W}$ as $s_i = f_{A_i}(w_i)$.

Then the convolution weights $w$ are modulated as follows. ($w$ here on refers to weights not intermediate latent space, we are sticking to the same notation as the paper.)

$$w'_{i, j, k} = s_i \cdot w_{i, j, k}$$ Then it’s demodulated by normalizing, $$w''_{i,j,k} = \frac{w'_{i,j,k}}{\sqrt{\sum_{i,k}{w'_{i, j, k}}^2 + \epsilon}}$$ where $i$ is the input channel, $j$ is the output channel, and $k$ is the kernel index.

Path Length Regularization

Path length regularization encourages a fixed-size step in $\mathcal{W}$ to result in a non-zero, fixed-magnitude change in the generated image.

No Progressive Growing

StyleGAN2 uses residual connections (with down-sampling) in the discriminator and skip connections in the generator with up-sampling (the RGB outputs from each layer are added - no residual connections in feature maps). They show that with experiments that the contribution of low-resolution layers is higher at beginning of the training and then high-resolution layers take over.

Mapping Network

This is an MLP with 8 linear layers. The mapping network maps the latent vector $z \in \mathcal{W}$ to an intermediate latent space $w \in \mathcal{W}$. $\mathcal{W}$ space will be disentangled from the image space where the factors of variation become more linear.

• features is the number of features in $z$ and $w$
• n_layers is the number of layers in the mapping network.

StyleGAN2 Generator

$A$ denotes a linear layer. $B$ denotes a broadcast and scaling operation (noise is a single channel). toRGB also has a style modulation which is not shown in the diagram to keep it simple.

The generator starts with a learned constant. Then it has a series of blocks. The feature map resolution is doubled at each block Each block outputs an RGB image and they are scaled up and summed to get the final RGB image.

• log_resolution is the $\log_2$ of image resolution
• d_latent is the dimensionality of $w$
• n_features number of features in the convolution layer at the highest resolution (final block)
• max_features maximum number of features in any generator block

• w is $w$. In order to mix-styles (use different $w$ for different layers), we provide a separate $w$ for each generator block. It has shape [n_blocks, batch_size, d_latent].
• input_noise is the noise for each block. It’s a list of pairs of noise sensors because each block (except the initial) has two noise inputs after each convolution layer (see the diagram).

Generator Block

$A$ denotes a linear layer. $B$ denotes a broadcast and scaling operation (noise is a single channel). toRGB also has a style modulation which is not shown in the diagram to keep it simple.

The generator block consists of two style blocks ($3 \times 3$ convolutions with style modulation) and an RGB output.

• d_latent is the dimensionality of $w$
• in_features is the number of features in the input feature map
• out_features is the number of features in the output feature map

• x is the input feature map of shape [batch_size, in_features, height, width]
• w is $w$ with shape [batch_size, d_latent]
• noise is a tuple of two noise tensors of shape [batch_size, 1, height, width]

Style Block

$A$ denotes a linear layer. $B$ denotes a broadcast and scaling operation (noise is single channel).

Style block has a weight modulation convolution layer.

• d_latent is the dimensionality of $w$
• in_features is the number of features in the input feature map
• out_features is the number of features in the output feature map

• x is the input feature map of shape [batch_size, in_features, height, width]
• w is $w$ with shape [batch_size, d_latent]
• noise is a tensor of shape [batch_size, 1, height, width]

To RGB

$A$ denotes a linear layer.

Generates an RGB image from a feature map using $1 \times 1$ convolution.

• d_latent is the dimensionality of $w$
• features is the number of features in the feature map

• x is the input feature map of shape [batch_size, in_features, height, width]
• w is $w$ with shape [batch_size, d_latent]

Convolution with Weight Modulation and Demodulation

This layer scales the convolution weights by the style vector and demodulates by normalizing it.

• in_features is the number of features in the input feature map
• out_features is the number of features in the output feature map
• kernel_size is the size of the convolution kernel
• demodulate is flag whether to normalize weights by its standard deviation
• eps is the $\epsilon$ for normalizing

• x is the input feature map of shape [batch_size, in_features, height, width]
• s is style based scaling tensor of shape [batch_size, in_features]

StyleGAN 2 Discriminator

Discriminator first transforms the image to a feature map of the same resolution and then runs it through a series of blocks with residual connections. The resolution is down-sampled by $2 \times$ at each block while doubling the number of features.

• log_resolution is the $\log_2$ of image resolution
• n_features number of features in the convolution layer at the highest resolution (first block)
• max_features maximum number of features in any generator block

• x is the input image of shape [batch_size, 3, height, width]

Discriminator Block

Discriminator block consists of two $3 \times 3$ convolutions with a residual connection.

• in_features is the number of features in the input feature map
• out_features is the number of features in the output feature map

Mini-batch Standard Deviation

Mini-batch standard deviation calculates the standard deviation across a mini-batch (or a subgroups within the mini-batch) for each feature in the feature map. Then it takes the mean of all the standard deviations and appends it to the feature map as one extra feature.

• group_size is the number of samples to calculate standard deviation across.

• x is the feature map

Down-sample

The down-sample operation smoothens each feature channel and scale $2 \times$ using bilinear interpolation. This is based on the paper Making Convolutional Networks Shift-Invariant Again.

Up-sample

The up-sample operation scales the image up by $2 \times$ and smoothens each feature channel. This is based on the paper Making Convolutional Networks Shift-Invariant Again.

Smoothing Layer

This layer blurs each channel

Learning-rate Equalized Linear Layer

This uses learning-rate equalized weights for a linear layer.

• in_features is the number of features in the input feature map
• out_features is the number of features in the output feature map
• bias is the bias initialization constant

Learning-rate Equalized 2D Convolution Layer

This uses learning-rate equalized weights for a convolution layer.

• in_features is the number of features in the input feature map
• out_features is the number of features in the output feature map
• kernel_size is the size of the convolution kernel
• padding is the padding to be added on both sides of each size dimension

Learning-rate Equalized Weights Parameter

This is based on equalized learning rate introduced in the Progressive GAN paper. Instead of initializing weights at $\mathcal{N}(0,c)$ they initialize weights to $\mathcal{N}(0, 1)$ and then multiply them by $c$ when using it. $$w_i = c \hat{w}_i$$

The gradients on stored parameters $\hat{w}$ get multiplied by $c$ but this doesn’t have an affect since optimizers such as Adam normalize them by a running mean of the squared gradients.

The optimizer updates on $\hat{w}$ are proportionate to the learning rate $\lambda$. But the effective weights $w$ get updated proportionately to $c \lambda$. Without equalized learning rate, the effective weights will get updated proportionately to just $\lambda$.

So we are effectively scaling the learning rate by $c$ for these weight parameters.

• shape is the shape of the weight parameter

Gradient Penalty

This is the $R_1$ regularization penality from the paper Which Training Methods for GANs do actually Converge?.

$$R_1(\psi) = \frac{\gamma}{2} \mathbb{E}_{p_\mathcal{D}(x)} \Big[\Vert \nabla_x D_\psi(x)^2 \Vert\Big]$$

That is we try to reduce the L2 norm of gradients of the discriminator with respect to images, for real images ($P_\mathcal{D}$).

• x is $x \sim \mathcal{D}$
• d is $D(x)$

Path Length Penalty

This regularization encourages a fixed-size step in $w$ to result in a fixed-magnitude change in the image.

$$\mathbb{E}_{w \sim f(z), y \sim \mathcal{N}(0, \mathbf{I})} \Big(\Vert \mathbf{J}^\top_{w} y \Vert_2 - a \Big)^2$$

where $\mathbf{J}_w$ is the Jacobian $\mathbf{J}_w = \frac{\partial g}{\partial w}$, $w$ are sampled from $w \in \mathcal{W}$ from the mapping network, and $y$ are images with noise $\mathcal{N}(0, \mathbf{I})$.

$a$ is the exponential moving average of $\Vert \mathbf{J}^\top_{w} y \Vert_2$ as the training progresses.

$\mathbf{J}^\top_{w} y$ is calculated without explicitly calculating the Jacobian using $$\mathbf{J}^\top_{w} y = \nabla_w \big(g(w) \cdot y \big)$$

• beta is the constant $\beta$ used to calculate the exponential moving average $a$

• w is the batch of $w$ of shape [batch_size, d_latent]
• x are the generated images of shape [batch_size, 3, height, width]