|
|
|
"""Contains the implementation of discriminator described in PGGAN. |
|
|
|
Paper: https://arxiv.org/pdf/1710.10196.pdf |
|
|
|
Official TensorFlow implementation: |
|
https://github.com/tkarras/progressive_growing_of_gans |
|
""" |
|
|
|
import numpy as np |
|
|
|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
|
|
__all__ = ['PGGANDiscriminator'] |
|
|
|
|
|
_RESOLUTIONS_ALLOWED = [8, 16, 32, 64, 128, 256, 512, 1024] |
|
|
|
|
|
_WSCALE_GAIN = np.sqrt(2.0) |
|
|
|
|
|
|
|
class PGGANDiscriminator(nn.Module): |
|
"""Defines the discriminator network in PGGAN. |
|
|
|
NOTE: The discriminator takes images with `RGB` channel order and pixel |
|
range [-1, 1] as inputs. |
|
|
|
Settings for the network: |
|
|
|
(1) resolution: The resolution of the input image. |
|
(2) init_res: Smallest resolution of the convolutional backbone. |
|
(default: 4) |
|
(3) image_channels: Number of channels of the input image. (default: 3) |
|
(4) label_dim: Dimension of the additional label for conditional generation. |
|
In one-hot conditioning case, it is equal to the number of classes. If |
|
set to 0, conditioning training will be disabled. (default: 0) |
|
(5) fused_scale: Whether to fused `conv2d` and `downsample` together, |
|
resulting in `conv2d` with strides. (default: False) |
|
(6) use_wscale: Whether to use weight scaling. (default: True) |
|
(7) wscale_gain: The factor to control weight scaling. (default: sqrt(2.0)) |
|
(8) mbstd_groups: Group size for the minibatch standard deviation layer. |
|
`0` means disable. (default: 16) |
|
(9) fmaps_base: Factor to control number of feature maps for each layer. |
|
(default: 16 << 10) |
|
(10) fmaps_max: Maximum number of feature maps in each layer. (default: 512) |
|
(11) eps: A small value to avoid divide overflow. (default: 1e-8) |
|
""" |
|
|
|
def __init__(self, |
|
resolution, |
|
init_res=4, |
|
image_channels=3, |
|
label_dim=0, |
|
fused_scale=False, |
|
use_wscale=True, |
|
wscale_gain=np.sqrt(2.0), |
|
mbstd_groups=16, |
|
fmaps_base=16 << 10, |
|
fmaps_max=512, |
|
eps=1e-8): |
|
"""Initializes with basic settings. |
|
|
|
Raises: |
|
ValueError: If the `resolution` is not supported. |
|
""" |
|
super().__init__() |
|
|
|
if resolution not in _RESOLUTIONS_ALLOWED: |
|
raise ValueError(f'Invalid resolution: `{resolution}`!\n' |
|
f'Resolutions allowed: {_RESOLUTIONS_ALLOWED}.') |
|
|
|
self.init_res = init_res |
|
self.init_res_log2 = int(np.log2(self.init_res)) |
|
self.resolution = resolution |
|
self.final_res_log2 = int(np.log2(self.resolution)) |
|
self.image_channels = image_channels |
|
self.label_dim = label_dim |
|
self.fused_scale = fused_scale |
|
self.use_wscale = use_wscale |
|
self.wscale_gain = wscale_gain |
|
self.mbstd_groups = mbstd_groups |
|
self.fmaps_base = fmaps_base |
|
self.fmaps_max = fmaps_max |
|
self.eps = eps |
|
|
|
|
|
self.register_buffer('lod', torch.zeros(())) |
|
self.pth_to_tf_var_mapping = {'lod': 'lod'} |
|
|
|
for res_log2 in range(self.final_res_log2, self.init_res_log2 - 1, -1): |
|
res = 2 ** res_log2 |
|
in_channels = self.get_nf(res) |
|
out_channels = self.get_nf(res // 2) |
|
block_idx = self.final_res_log2 - res_log2 |
|
|
|
|
|
self.add_module( |
|
f'input{block_idx}', |
|
ConvLayer(in_channels=self.image_channels, |
|
out_channels=in_channels, |
|
kernel_size=1, |
|
add_bias=True, |
|
downsample=False, |
|
fused_scale=False, |
|
use_wscale=use_wscale, |
|
wscale_gain=wscale_gain, |
|
activation_type='lrelu')) |
|
self.pth_to_tf_var_mapping[f'input{block_idx}.weight'] = ( |
|
f'FromRGB_lod{block_idx}/weight') |
|
self.pth_to_tf_var_mapping[f'input{block_idx}.bias'] = ( |
|
f'FromRGB_lod{block_idx}/bias') |
|
|
|
|
|
if res != self.init_res: |
|
self.add_module( |
|
f'layer{2 * block_idx}', |
|
ConvLayer(in_channels=in_channels, |
|
out_channels=in_channels, |
|
kernel_size=3, |
|
add_bias=True, |
|
downsample=False, |
|
fused_scale=False, |
|
use_wscale=use_wscale, |
|
wscale_gain=wscale_gain, |
|
activation_type='lrelu')) |
|
tf_layer0_name = 'Conv0' |
|
self.add_module( |
|
f'layer{2 * block_idx + 1}', |
|
ConvLayer(in_channels=in_channels, |
|
out_channels=out_channels, |
|
kernel_size=3, |
|
add_bias=True, |
|
downsample=True, |
|
fused_scale=fused_scale, |
|
use_wscale=use_wscale, |
|
wscale_gain=wscale_gain, |
|
activation_type='lrelu')) |
|
tf_layer1_name = 'Conv1_down' if fused_scale else 'Conv1' |
|
|
|
|
|
else: |
|
self.mbstd = MiniBatchSTDLayer(groups=mbstd_groups, eps=eps) |
|
self.add_module( |
|
f'layer{2 * block_idx}', |
|
ConvLayer( |
|
in_channels=in_channels + 1, |
|
out_channels=in_channels, |
|
kernel_size=3, |
|
add_bias=True, |
|
downsample=False, |
|
fused_scale=False, |
|
use_wscale=use_wscale, |
|
wscale_gain=wscale_gain, |
|
activation_type='lrelu')) |
|
tf_layer0_name = 'Conv' |
|
self.add_module( |
|
f'layer{2 * block_idx + 1}', |
|
DenseLayer(in_channels=in_channels * res * res, |
|
out_channels=out_channels, |
|
add_bias=True, |
|
use_wscale=use_wscale, |
|
wscale_gain=wscale_gain, |
|
activation_type='lrelu')) |
|
tf_layer1_name = 'Dense0' |
|
|
|
self.pth_to_tf_var_mapping[f'layer{2 * block_idx}.weight'] = ( |
|
f'{res}x{res}/{tf_layer0_name}/weight') |
|
self.pth_to_tf_var_mapping[f'layer{2 * block_idx}.bias'] = ( |
|
f'{res}x{res}/{tf_layer0_name}/bias') |
|
self.pth_to_tf_var_mapping[f'layer{2 * block_idx + 1}.weight'] = ( |
|
f'{res}x{res}/{tf_layer1_name}/weight') |
|
self.pth_to_tf_var_mapping[f'layer{2 * block_idx + 1}.bias'] = ( |
|
f'{res}x{res}/{tf_layer1_name}/bias') |
|
|
|
|
|
self.output = DenseLayer(in_channels=out_channels, |
|
out_channels=1 + self.label_dim, |
|
add_bias=True, |
|
use_wscale=self.use_wscale, |
|
wscale_gain=1.0, |
|
activation_type='linear') |
|
self.pth_to_tf_var_mapping['output.weight'] = ( |
|
f'{res}x{res}/Dense1/weight') |
|
self.pth_to_tf_var_mapping['output.bias'] = ( |
|
f'{res}x{res}/Dense1/bias') |
|
|
|
def get_nf(self, res): |
|
"""Gets number of feature maps according to the given resolution.""" |
|
return min(self.fmaps_base // res, self.fmaps_max) |
|
|
|
def forward(self, image, lod=None): |
|
expected_shape = (self.image_channels, self.resolution, self.resolution) |
|
if image.ndim != 4 or image.shape[1:] != expected_shape: |
|
raise ValueError(f'The input tensor should be with shape ' |
|
f'[batch_size, channel, height, width], where ' |
|
f'`channel` equals to {self.image_channels}, ' |
|
f'`height`, `width` equal to {self.resolution}!\n' |
|
f'But `{image.shape}` is received!') |
|
|
|
lod = self.lod.item() if lod is None else lod |
|
if lod + self.init_res_log2 > self.final_res_log2: |
|
raise ValueError(f'Maximum level-of-details (lod) is ' |
|
f'{self.final_res_log2 - self.init_res_log2}, ' |
|
f'but `{lod}` is received!') |
|
|
|
lod = self.lod.item() |
|
for res_log2 in range(self.final_res_log2, self.init_res_log2 - 1, -1): |
|
block_idx = current_lod = self.final_res_log2 - res_log2 |
|
if current_lod <= lod < current_lod + 1: |
|
x = getattr(self, f'input{block_idx}')(image) |
|
elif current_lod - 1 < lod < current_lod: |
|
alpha = lod - np.floor(lod) |
|
y = getattr(self, f'input{block_idx}')(image) |
|
x = y * alpha + x * (1 - alpha) |
|
if lod < current_lod + 1: |
|
if res_log2 == self.init_res_log2: |
|
x = self.mbstd(x) |
|
x = getattr(self, f'layer{2 * block_idx}')(x) |
|
x = getattr(self, f'layer{2 * block_idx + 1}')(x) |
|
if lod > current_lod: |
|
image = F.avg_pool2d( |
|
image, kernel_size=2, stride=2, padding=0) |
|
x = self.output(x) |
|
|
|
return {'score': x} |
|
|
|
|
|
class MiniBatchSTDLayer(nn.Module): |
|
"""Implements the minibatch standard deviation layer.""" |
|
|
|
def __init__(self, groups, eps): |
|
super().__init__() |
|
self.groups = groups |
|
self.eps = eps |
|
|
|
def extra_repr(self): |
|
return f'groups={self.groups}, epsilon={self.eps}' |
|
|
|
def forward(self, x): |
|
if self.groups <= 1: |
|
return x |
|
|
|
N, C, H, W = x.shape |
|
G = min(self.groups, N) |
|
|
|
y = x.reshape(G, -1, C, H, W) |
|
y = y - y.mean(dim=0) |
|
y = y.square().mean(dim=0) |
|
y = (y + self.eps).sqrt() |
|
y = y.mean(dim=(1, 2, 3), keepdim=True) |
|
y = y.repeat(G, 1, H, W) |
|
x = torch.cat([x, y], dim=1) |
|
|
|
return x |
|
|
|
|
|
class DownsamplingLayer(nn.Module): |
|
"""Implements the downsampling layer. |
|
|
|
Basically, this layer can be used to downsample feature maps with average |
|
pooling. |
|
""" |
|
|
|
def __init__(self, scale_factor): |
|
super().__init__() |
|
self.scale_factor = scale_factor |
|
|
|
def extra_repr(self): |
|
return f'factor={self.scale_factor}' |
|
|
|
def forward(self, x): |
|
if self.scale_factor <= 1: |
|
return x |
|
return F.avg_pool2d(x, |
|
kernel_size=self.scale_factor, |
|
stride=self.scale_factor, |
|
padding=0) |
|
|
|
|
|
class ConvLayer(nn.Module): |
|
"""Implements the convolutional layer. |
|
|
|
Basically, this layer executes convolution, activation, and downsampling (if |
|
needed) in sequence. |
|
""" |
|
|
|
def __init__(self, |
|
in_channels, |
|
out_channels, |
|
kernel_size, |
|
add_bias, |
|
downsample, |
|
fused_scale, |
|
use_wscale, |
|
wscale_gain, |
|
activation_type): |
|
"""Initializes with layer settings. |
|
|
|
Args: |
|
in_channels: Number of channels of the input tensor. |
|
out_channels: Number of channels of the output tensor. |
|
kernel_size: Size of the convolutional kernels. |
|
add_bias: Whether to add bias onto the convolutional result. |
|
downsample: Whether to downsample the result after convolution. |
|
fused_scale: Whether to fused `conv2d` and `downsample` together, |
|
resulting in `conv2d` with strides. |
|
use_wscale: Whether to use weight scaling. |
|
wscale_gain: Gain factor for weight scaling. |
|
activation_type: Type of activation. |
|
""" |
|
super().__init__() |
|
self.in_channels = in_channels |
|
self.out_channels = out_channels |
|
self.kernel_size = kernel_size |
|
self.add_bias = add_bias |
|
self.downsample = downsample |
|
self.fused_scale = fused_scale |
|
self.use_wscale = use_wscale |
|
self.wscale_gain = wscale_gain |
|
self.activation_type = activation_type |
|
|
|
if downsample and not fused_scale: |
|
self.down = DownsamplingLayer(scale_factor=2) |
|
else: |
|
self.down = nn.Identity() |
|
|
|
if downsample and fused_scale: |
|
self.use_stride = True |
|
self.stride = 2 |
|
self.padding = 1 |
|
else: |
|
self.use_stride = False |
|
self.stride = 1 |
|
self.padding = kernel_size // 2 |
|
|
|
weight_shape = (out_channels, in_channels, kernel_size, kernel_size) |
|
fan_in = kernel_size * kernel_size * in_channels |
|
wscale = wscale_gain / np.sqrt(fan_in) |
|
if use_wscale: |
|
self.weight = nn.Parameter(torch.randn(*weight_shape)) |
|
self.wscale = wscale |
|
else: |
|
self.weight = nn.Parameter(torch.randn(*weight_shape) * wscale) |
|
self.wscale = 1.0 |
|
|
|
if add_bias: |
|
self.bias = nn.Parameter(torch.zeros(out_channels)) |
|
else: |
|
self.bias = None |
|
|
|
assert activation_type in ['linear', 'relu', 'lrelu'] |
|
|
|
def extra_repr(self): |
|
return (f'in_ch={self.in_channels}, ' |
|
f'out_ch={self.out_channels}, ' |
|
f'ksize={self.kernel_size}, ' |
|
f'wscale_gain={self.wscale_gain:.3f}, ' |
|
f'bias={self.add_bias}, ' |
|
f'downsample={self.scale_factor}, ' |
|
f'fused_scale={self.fused_scale}, ' |
|
f'act={self.activation_type}') |
|
|
|
def forward(self, x): |
|
weight = self.weight |
|
if self.wscale != 1.0: |
|
weight = weight * self.wscale |
|
|
|
if self.use_stride: |
|
weight = F.pad(weight, (1, 1, 1, 1, 0, 0, 0, 0), 'constant', 0.0) |
|
weight = (weight[:, :, 1:, 1:] + weight[:, :, :-1, 1:] + |
|
weight[:, :, 1:, :-1] + weight[:, :, :-1, :-1]) * 0.25 |
|
x = F.conv2d(x, |
|
weight=weight, |
|
bias=self.bias, |
|
stride=self.stride, |
|
padding=self.padding) |
|
|
|
if self.activation_type == 'linear': |
|
pass |
|
elif self.activation_type == 'relu': |
|
x = F.relu(x, inplace=True) |
|
elif self.activation_type == 'lrelu': |
|
x = F.leaky_relu(x, negative_slope=0.2, inplace=True) |
|
else: |
|
raise NotImplementedError(f'Not implemented activation type ' |
|
f'`{self.activation_type}`!') |
|
x = self.down(x) |
|
|
|
return x |
|
|
|
|
|
class DenseLayer(nn.Module): |
|
"""Implements the dense layer.""" |
|
|
|
def __init__(self, |
|
in_channels, |
|
out_channels, |
|
add_bias, |
|
use_wscale, |
|
wscale_gain, |
|
activation_type): |
|
"""Initializes with layer settings. |
|
|
|
Args: |
|
in_channels: Number of channels of the input tensor. |
|
out_channels: Number of channels of the output tensor. |
|
add_bias: Whether to add bias onto the fully-connected result. |
|
use_wscale: Whether to use weight scaling. |
|
wscale_gain: Gain factor for weight scaling. |
|
activation_type: Type of activation. |
|
|
|
Raises: |
|
NotImplementedError: If the `activation_type` is not supported. |
|
""" |
|
super().__init__() |
|
self.in_channels = in_channels |
|
self.out_channels = out_channels |
|
self.add_bias = add_bias |
|
self.use_wscale = use_wscale |
|
self.wscale_gain = wscale_gain |
|
self.activation_type = activation_type |
|
|
|
weight_shape = (out_channels, in_channels) |
|
wscale = wscale_gain / np.sqrt(in_channels) |
|
if use_wscale: |
|
self.weight = nn.Parameter(torch.randn(*weight_shape)) |
|
self.wscale = wscale |
|
else: |
|
self.weight = nn.Parameter(torch.randn(*weight_shape) * wscale) |
|
self.wscale = 1.0 |
|
|
|
if add_bias: |
|
self.bias = nn.Parameter(torch.zeros(out_channels)) |
|
else: |
|
self.bias = None |
|
|
|
assert activation_type in ['linear', 'relu', 'lrelu'] |
|
|
|
def forward(self, x): |
|
if x.ndim != 2: |
|
x = x.flatten(start_dim=1) |
|
|
|
weight = self.weight |
|
if self.wscale != 1.0: |
|
weight = weight * self.wscale |
|
|
|
x = F.linear(x, weight=weight, bias=self.bias) |
|
|
|
if self.activation_type == 'linear': |
|
pass |
|
elif self.activation_type == 'relu': |
|
x = F.relu(x, inplace=True) |
|
elif self.activation_type == 'lrelu': |
|
x = F.leaky_relu(x, negative_slope=0.2, inplace=True) |
|
else: |
|
raise NotImplementedError(f'Not implemented activation type ' |
|
f'`{self.activation_type}`!') |
|
|
|
return x |
|
|
|
|
|
|
