Spaces:
Build error
Build error
# python3.7 | |
"""Contains the implementation of generator described in StyleGAN. | |
Paper: https://arxiv.org/pdf/1812.04948.pdf | |
Official TensorFlow implementation: https://github.com/NVlabs/stylegan | |
""" | |
import numpy as np | |
import torch | |
import torch.nn as nn | |
import torch.nn.functional as F | |
from .sync_op import all_gather | |
__all__ = ['StyleGANGenerator'] | |
# Resolutions allowed. | |
_RESOLUTIONS_ALLOWED = [8, 16, 32, 64, 128, 256, 512, 1024] | |
# Initial resolution. | |
_INIT_RES = 4 | |
# Fused-scale options allowed. | |
_FUSED_SCALE_ALLOWED = [True, False, 'auto'] | |
# Minimal resolution for `auto` fused-scale strategy. | |
_AUTO_FUSED_SCALE_MIN_RES = 128 | |
# Default gain factor for weight scaling. | |
_WSCALE_GAIN = np.sqrt(2.0) | |
_STYLEMOD_WSCALE_GAIN = 1.0 | |
class StyleGANGenerator(nn.Module): | |
"""Defines the generator network in StyleGAN. | |
NOTE: The synthesized images are with `RGB` channel order and pixel range | |
[-1, 1]. | |
Settings for the mapping network: | |
(1) z_space_dim: Dimension of the input latent space, Z. (default: 512) | |
(2) w_space_dim: Dimension of the outout latent space, W. (default: 512) | |
(3) label_size: Size of the additional label for conditional generation. | |
(default: 0) | |
(4)mapping_layers: Number of layers of the mapping network. (default: 8) | |
(5) mapping_fmaps: Number of hidden channels of the mapping network. | |
(default: 512) | |
(6) mapping_lr_mul: Learning rate multiplier for the mapping network. | |
(default: 0.01) | |
(7) repeat_w: Repeat w-code for different layers. | |
Settings for the synthesis network: | |
(1) resolution: The resolution of the output image. | |
(2) image_channels: Number of channels of the output image. (default: 3) | |
(3) final_tanh: Whether to use `tanh` to control the final pixel range. | |
(default: False) | |
(4) const_input: Whether to use a constant in the first convolutional layer. | |
(default: True) | |
(5) fused_scale: Whether to fused `upsample` and `conv2d` together, | |
resulting in `conv2d_transpose`. (default: `auto`) | |
(6) use_wscale: Whether to use weight scaling. (default: True) | |
(7) noise_type: Type of noise added to the convolutional results at each | |
layer. (default: `spatial`) | |
(8) fmaps_base: Factor to control number of feature maps for each layer. | |
(default: 16 << 10) | |
(9) fmaps_max: Maximum number of feature maps in each layer. (default: 512) | |
""" | |
def __init__(self, | |
resolution, | |
z_space_dim=512, | |
w_space_dim=512, | |
label_size=0, | |
mapping_layers=8, | |
mapping_fmaps=512, | |
mapping_lr_mul=0.01, | |
repeat_w=True, | |
image_channels=3, | |
final_tanh=False, | |
const_input=True, | |
fused_scale='auto', | |
use_wscale=True, | |
noise_type='spatial', | |
fmaps_base=16 << 10, | |
fmaps_max=512): | |
"""Initializes with basic settings. | |
Raises: | |
ValueError: If the `resolution` is not supported, or `fused_scale` | |
is not supported. | |
""" | |
super().__init__() | |
if resolution not in _RESOLUTIONS_ALLOWED: | |
raise ValueError(f'Invalid resolution: `{resolution}`!\n' | |
f'Resolutions allowed: {_RESOLUTIONS_ALLOWED}.') | |
if fused_scale not in _FUSED_SCALE_ALLOWED: | |
raise ValueError(f'Invalid fused-scale option: `{fused_scale}`!\n' | |
f'Options allowed: {_FUSED_SCALE_ALLOWED}.') | |
self.init_res = _INIT_RES | |
self.resolution = resolution | |
self.z_space_dim = z_space_dim | |
self.w_space_dim = w_space_dim | |
self.label_size = label_size | |
self.mapping_layers = mapping_layers | |
self.mapping_fmaps = mapping_fmaps | |
self.mapping_lr_mul = mapping_lr_mul | |
self.repeat_w = repeat_w | |
self.image_channels = image_channels | |
self.final_tanh = final_tanh | |
self.const_input = const_input | |
self.fused_scale = fused_scale | |
self.use_wscale = use_wscale | |
self.noise_type = noise_type | |
self.fmaps_base = fmaps_base | |
self.fmaps_max = fmaps_max | |
self.num_layers = int(np.log2(self.resolution // self.init_res * 2)) * 2 | |
if self.repeat_w: | |
self.mapping_space_dim = self.w_space_dim | |
else: | |
self.mapping_space_dim = self.w_space_dim * self.num_layers | |
self.mapping = MappingModule(input_space_dim=self.z_space_dim, | |
hidden_space_dim=self.mapping_fmaps, | |
final_space_dim=self.mapping_space_dim, | |
label_size=self.label_size, | |
num_layers=self.mapping_layers, | |
use_wscale=self.use_wscale, | |
lr_mul=self.mapping_lr_mul) | |
self.truncation = TruncationModule(w_space_dim=self.w_space_dim, | |
num_layers=self.num_layers, | |
repeat_w=self.repeat_w) | |
self.synthesis = SynthesisModule(resolution=self.resolution, | |
init_resolution=self.init_res, | |
w_space_dim=self.w_space_dim, | |
image_channels=self.image_channels, | |
final_tanh=self.final_tanh, | |
const_input=self.const_input, | |
fused_scale=self.fused_scale, | |
use_wscale=self.use_wscale, | |
noise_type=self.noise_type, | |
fmaps_base=self.fmaps_base, | |
fmaps_max=self.fmaps_max) | |
self.pth_to_tf_var_mapping = {} | |
for key, val in self.mapping.pth_to_tf_var_mapping.items(): | |
self.pth_to_tf_var_mapping[f'mapping.{key}'] = val | |
for key, val in self.truncation.pth_to_tf_var_mapping.items(): | |
self.pth_to_tf_var_mapping[f'truncation.{key}'] = val | |
for key, val in self.synthesis.pth_to_tf_var_mapping.items(): | |
self.pth_to_tf_var_mapping[f'synthesis.{key}'] = val | |
def set_space_of_latent(self, space_of_latent='w'): | |
"""Sets the space to which the latent code belong. | |
This function is particually used for choosing how to inject the latent | |
code into the convolutional layers. The original generator will take a | |
W-Space code and apply it for style modulation after an affine | |
transformation. But, sometimes, it may need to directly feed an already | |
affine-transformed code into the convolutional layer, e.g., when | |
training an encoder for GAN inversion. We term the transformed space as | |
Style Space (or Y-Space). This function is designed to tell the | |
convolutional layers how to use the input code. | |
Args: | |
space_of_latent: The space to which the latent code belong. Case | |
insensitive. (default: 'w') | |
""" | |
for module in self.modules(): | |
if isinstance(module, StyleModLayer): | |
setattr(module, 'space_of_latent', space_of_latent) | |
def forward(self, | |
z, | |
label=None, | |
lod=None, | |
w_moving_decay=0.995, | |
style_mixing_prob=0.9, | |
trunc_psi=None, | |
trunc_layers=None, | |
randomize_noise=False, | |
**_unused_kwargs): | |
mapping_results = self.mapping(z, label) | |
w = mapping_results['w'] | |
if self.training and w_moving_decay < 1: | |
batch_w_avg = all_gather(w).mean(dim=0) | |
self.truncation.w_avg.copy_( | |
self.truncation.w_avg * w_moving_decay + | |
batch_w_avg * (1 - w_moving_decay)) | |
if self.training and style_mixing_prob > 0: | |
new_z = torch.randn_like(z) | |
new_w = self.mapping(new_z, label)['w'] | |
lod = self.synthesis.lod.cpu().tolist() if lod is None else lod | |
current_layers = self.num_layers - int(lod) * 2 | |
if np.random.uniform() < style_mixing_prob: | |
mixing_cutoff = np.random.randint(1, current_layers) | |
w = self.truncation(w) | |
new_w = self.truncation(new_w) | |
w[:, mixing_cutoff:] = new_w[:, mixing_cutoff:] | |
wp = self.truncation(w, trunc_psi, trunc_layers) | |
synthesis_results = self.synthesis(wp, lod, randomize_noise) | |
return {**mapping_results, **synthesis_results} | |
class MappingModule(nn.Module): | |
"""Implements the latent space mapping module. | |
Basically, this module executes several dense layers in sequence. | |
""" | |
def __init__(self, | |
input_space_dim=512, | |
hidden_space_dim=512, | |
final_space_dim=512, | |
label_size=0, | |
num_layers=8, | |
normalize_input=True, | |
use_wscale=True, | |
lr_mul=0.01): | |
super().__init__() | |
self.input_space_dim = input_space_dim | |
self.hidden_space_dim = hidden_space_dim | |
self.final_space_dim = final_space_dim | |
self.label_size = label_size | |
self.num_layers = num_layers | |
self.normalize_input = normalize_input | |
self.use_wscale = use_wscale | |
self.lr_mul = lr_mul | |
self.norm = PixelNormLayer() if self.normalize_input else nn.Identity() | |
self.pth_to_tf_var_mapping = {} | |
for i in range(num_layers): | |
dim_mul = 2 if label_size else 1 | |
in_channels = (input_space_dim * dim_mul if i == 0 else | |
hidden_space_dim) | |
out_channels = (final_space_dim if i == (num_layers - 1) else | |
hidden_space_dim) | |
self.add_module(f'dense{i}', | |
DenseBlock(in_channels=in_channels, | |
out_channels=out_channels, | |
use_wscale=self.use_wscale, | |
lr_mul=self.lr_mul)) | |
self.pth_to_tf_var_mapping[f'dense{i}.weight'] = f'Dense{i}/weight' | |
self.pth_to_tf_var_mapping[f'dense{i}.bias'] = f'Dense{i}/bias' | |
if label_size: | |
self.label_weight = nn.Parameter( | |
torch.randn(label_size, input_space_dim)) | |
self.pth_to_tf_var_mapping[f'label_weight'] = f'LabelConcat/weight' | |
def forward(self, z, label=None): | |
if z.ndim != 2 or z.shape[1] != self.input_space_dim: | |
raise ValueError(f'Input latent code should be with shape ' | |
f'[batch_size, input_dim], where ' | |
f'`input_dim` equals to {self.input_space_dim}!\n' | |
f'But `{z.shape}` is received!') | |
if self.label_size: | |
if label is None: | |
raise ValueError(f'Model requires an additional label ' | |
f'(with size {self.label_size}) as input, ' | |
f'but no label is received!') | |
if label.ndim != 2 or label.shape != (z.shape[0], self.label_size): | |
raise ValueError(f'Input label should be with shape ' | |
f'[batch_size, label_size], where ' | |
f'`batch_size` equals to that of ' | |
f'latent codes ({z.shape[0]}) and ' | |
f'`label_size` equals to {self.label_size}!\n' | |
f'But `{label.shape}` is received!') | |
embedding = torch.matmul(label, self.label_weight) | |
z = torch.cat((z, embedding), dim=1) | |
z = self.norm(z) | |
w = z | |
for i in range(self.num_layers): | |
w = self.__getattr__(f'dense{i}')(w) | |
results = { | |
'z': z, | |
'label': label, | |
'w': w, | |
} | |
if self.label_size: | |
results['embedding'] = embedding | |
return results | |
class TruncationModule(nn.Module): | |
"""Implements the truncation module. | |
Truncation is executed as follows: | |
For layers in range [0, truncation_layers), the truncated w-code is computed | |
as | |
w_new = w_avg + (w - w_avg) * truncation_psi | |
To disable truncation, please set | |
(1) truncation_psi = 1.0 (None) OR | |
(2) truncation_layers = 0 (None) | |
NOTE: The returned tensor is layer-wise style codes. | |
""" | |
def __init__(self, w_space_dim, num_layers, repeat_w=True): | |
super().__init__() | |
self.num_layers = num_layers | |
self.w_space_dim = w_space_dim | |
self.repeat_w = repeat_w | |
if self.repeat_w: | |
self.register_buffer('w_avg', torch.zeros(w_space_dim)) | |
else: | |
self.register_buffer('w_avg', torch.zeros(num_layers * w_space_dim)) | |
self.pth_to_tf_var_mapping = {'w_avg': 'dlatent_avg'} | |
def forward(self, w, trunc_psi=None, trunc_layers=None): | |
if w.ndim == 2: | |
if self.repeat_w and w.shape[1] == self.w_space_dim: | |
w = w.view(-1, 1, self.w_space_dim) | |
wp = w.repeat(1, self.num_layers, 1) | |
else: | |
assert w.shape[1] == self.w_space_dim * self.num_layers | |
wp = w.view(-1, self.num_layers, self.w_space_dim) | |
else: | |
wp = w | |
assert wp.ndim == 3 | |
assert wp.shape[1:] == (self.num_layers, self.w_space_dim) | |
trunc_psi = 1.0 if trunc_psi is None else trunc_psi | |
trunc_layers = 0 if trunc_layers is None else trunc_layers | |
if trunc_psi < 1.0 and trunc_layers > 0: | |
layer_idx = np.arange(self.num_layers).reshape(1, -1, 1) | |
coefs = np.ones_like(layer_idx, dtype=np.float32) | |
coefs[layer_idx < trunc_layers] *= trunc_psi | |
coefs = torch.from_numpy(coefs).to(wp) | |
w_avg = self.w_avg.view(1, -1, self.w_space_dim) | |
wp = w_avg + (wp - w_avg) * coefs | |
return wp | |
class SynthesisModule(nn.Module): | |
"""Implements the image synthesis module. | |
Basically, this module executes several convolutional layers in sequence. | |
""" | |
def __init__(self, | |
resolution=1024, | |
init_resolution=4, | |
w_space_dim=512, | |
image_channels=3, | |
final_tanh=False, | |
const_input=True, | |
fused_scale='auto', | |
use_wscale=True, | |
noise_type='spatial', | |
fmaps_base=16 << 10, | |
fmaps_max=512): | |
super().__init__() | |
self.init_res = init_resolution | |
self.init_res_log2 = int(np.log2(self.init_res)) | |
self.resolution = resolution | |
self.final_res_log2 = int(np.log2(self.resolution)) | |
self.w_space_dim = w_space_dim | |
self.image_channels = image_channels | |
self.final_tanh = final_tanh | |
self.const_input = const_input | |
self.fused_scale = fused_scale | |
self.use_wscale = use_wscale | |
self.noise_type = noise_type | |
self.fmaps_base = fmaps_base | |
self.fmaps_max = fmaps_max | |
self.num_layers = (self.final_res_log2 - self.init_res_log2 + 1) * 2 | |
# Level of detail (used for progressive training). | |
self.register_buffer('lod', torch.zeros(())) | |
self.pth_to_tf_var_mapping = {'lod': 'lod'} | |
for res_log2 in range(self.init_res_log2, self.final_res_log2 + 1): | |
res = 2 ** res_log2 | |
block_idx = res_log2 - self.init_res_log2 | |
# First convolution layer for each resolution. | |
layer_name = f'layer{2 * block_idx}' | |
if res == self.init_res: | |
if self.const_input: | |
self.add_module(layer_name, | |
ConvBlock(in_channels=self.get_nf(res), | |
out_channels=self.get_nf(res), | |
resolution=self.init_res, | |
w_space_dim=self.w_space_dim, | |
position='const_init', | |
use_wscale=self.use_wscale, | |
noise_type=self.noise_type)) | |
tf_layer_name = 'Const' | |
self.pth_to_tf_var_mapping[f'{layer_name}.const'] = ( | |
f'{res}x{res}/{tf_layer_name}/const') | |
else: | |
self.add_module(layer_name, | |
ConvBlock(in_channels=self.w_space_dim, | |
out_channels=self.get_nf(res), | |
resolution=self.init_res, | |
w_space_dim=self.w_space_dim, | |
kernel_size=self.init_res, | |
padding=self.init_res - 1, | |
use_wscale=self.use_wscale, | |
noise_type=self.noise_type)) | |
tf_layer_name = 'Dense' | |
self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = ( | |
f'{res}x{res}/{tf_layer_name}/weight') | |
else: | |
if self.fused_scale == 'auto': | |
fused_scale = (res >= _AUTO_FUSED_SCALE_MIN_RES) | |
else: | |
fused_scale = self.fused_scale | |
self.add_module(layer_name, | |
ConvBlock(in_channels=self.get_nf(res // 2), | |
out_channels=self.get_nf(res), | |
resolution=res, | |
w_space_dim=self.w_space_dim, | |
upsample=True, | |
fused_scale=fused_scale, | |
use_wscale=self.use_wscale, | |
noise_type=self.noise_type)) | |
tf_layer_name = 'Conv0_up' | |
self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = ( | |
f'{res}x{res}/{tf_layer_name}/weight') | |
self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = ( | |
f'{res}x{res}/{tf_layer_name}/bias') | |
self.pth_to_tf_var_mapping[f'{layer_name}.style.weight'] = ( | |
f'{res}x{res}/{tf_layer_name}/StyleMod/weight') | |
self.pth_to_tf_var_mapping[f'{layer_name}.style.bias'] = ( | |
f'{res}x{res}/{tf_layer_name}/StyleMod/bias') | |
self.pth_to_tf_var_mapping[f'{layer_name}.apply_noise.weight'] = ( | |
f'{res}x{res}/{tf_layer_name}/Noise/weight') | |
self.pth_to_tf_var_mapping[f'{layer_name}.apply_noise.noise'] = ( | |
f'noise{2 * block_idx}') | |
# Second convolution layer for each resolution. | |
layer_name = f'layer{2 * block_idx + 1}' | |
self.add_module(layer_name, | |
ConvBlock(in_channels=self.get_nf(res), | |
out_channels=self.get_nf(res), | |
resolution=res, | |
w_space_dim=self.w_space_dim, | |
use_wscale=self.use_wscale, | |
noise_type=self.noise_type)) | |
tf_layer_name = 'Conv' if res == self.init_res else 'Conv1' | |
self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = ( | |
f'{res}x{res}/{tf_layer_name}/weight') | |
self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = ( | |
f'{res}x{res}/{tf_layer_name}/bias') | |
self.pth_to_tf_var_mapping[f'{layer_name}.style.weight'] = ( | |
f'{res}x{res}/{tf_layer_name}/StyleMod/weight') | |
self.pth_to_tf_var_mapping[f'{layer_name}.style.bias'] = ( | |
f'{res}x{res}/{tf_layer_name}/StyleMod/bias') | |
self.pth_to_tf_var_mapping[f'{layer_name}.apply_noise.weight'] = ( | |
f'{res}x{res}/{tf_layer_name}/Noise/weight') | |
self.pth_to_tf_var_mapping[f'{layer_name}.apply_noise.noise'] = ( | |
f'noise{2 * block_idx + 1}') | |
# Output convolution layer for each resolution. | |
self.add_module(f'output{block_idx}', | |
ConvBlock(in_channels=self.get_nf(res), | |
out_channels=self.image_channels, | |
resolution=res, | |
w_space_dim=self.w_space_dim, | |
position='last', | |
kernel_size=1, | |
padding=0, | |
use_wscale=self.use_wscale, | |
wscale_gain=1.0, | |
activation_type='linear')) | |
self.pth_to_tf_var_mapping[f'output{block_idx}.weight'] = ( | |
f'ToRGB_lod{self.final_res_log2 - res_log2}/weight') | |
self.pth_to_tf_var_mapping[f'output{block_idx}.bias'] = ( | |
f'ToRGB_lod{self.final_res_log2 - res_log2}/bias') | |
self.upsample = UpsamplingLayer() | |
self.final_activate = nn.Tanh() if final_tanh else nn.Identity() | |
def get_nf(self, res): | |
"""Gets number of feature maps according to current resolution.""" | |
return min(self.fmaps_base // res, self.fmaps_max) | |
def forward(self, wp, x=None, lod=None, randomize_noise=False, start=2, stop=None): | |
stop = self.final_res_log2 + 1 if stop is None else stop | |
lod = self.lod.cpu().tolist() if lod is None else lod | |
if lod + self.init_res_log2 > self.final_res_log2: | |
raise ValueError(f'Maximum level-of-detail (lod) is ' | |
f'{self.final_res_log2 - self.init_res_log2}, ' | |
f'but `{lod}` is received!') | |
results = {'wp': wp} | |
# for res_log2 in range(self.init_res_log2, self.final_res_log2 + 1): | |
for res_log2 in range(start, stop): | |
current_lod = self.final_res_log2 - res_log2 | |
if lod < current_lod + 1: | |
block_idx = res_log2 - self.init_res_log2 | |
if block_idx == 0: | |
if self.const_input: | |
x, style = self.layer0(None, wp[:, 0], randomize_noise) | |
else: | |
x = wp[:, 0].view(-1, self.w_space_dim, 1, 1) | |
x, style = self.layer0(x, wp[:, 0], randomize_noise) | |
else: | |
x, style = self.__getattr__(f'layer{2 * block_idx}')( | |
x, wp[:, 2 * block_idx]) | |
results[f'style{2 * block_idx:02d}'] = style | |
x, style = self.__getattr__(f'layer{2 * block_idx + 1}')( | |
x, wp[:, 2 * block_idx + 1]) | |
results[f'style{2 * block_idx + 1:02d}'] = style | |
if current_lod - 1 < lod <= current_lod: | |
image = self.__getattr__(f'output{block_idx}')(x, None) | |
elif current_lod < lod < current_lod + 1: | |
alpha = np.ceil(lod) - lod | |
image = (self.__getattr__(f'output{block_idx}')(x, None) * alpha | |
+ self.upsample(image) * (1 - alpha)) | |
elif lod >= current_lod + 1: | |
image = self.upsample(image) | |
results['image'] = self.final_activate(image) if res_log2 == self.final_res_log2 else None | |
results['x'] = x | |
return results | |
class PixelNormLayer(nn.Module): | |
"""Implements pixel-wise feature vector normalization layer.""" | |
def __init__(self, epsilon=1e-8): | |
super().__init__() | |
self.eps = epsilon | |
def forward(self, x): | |
norm = torch.sqrt(torch.mean(x ** 2, dim=1, keepdim=True) + self.eps) | |
return x / norm | |
class InstanceNormLayer(nn.Module): | |
"""Implements instance normalization layer.""" | |
def __init__(self, epsilon=1e-8): | |
super().__init__() | |
self.eps = epsilon | |
def forward(self, x): | |
if x.ndim != 4: | |
raise ValueError(f'The input tensor should be with shape ' | |
f'[batch_size, channel, height, width], ' | |
f'but `{x.shape}` is received!') | |
x = x - torch.mean(x, dim=[2, 3], keepdim=True) | |
norm = torch.sqrt( | |
torch.mean(x ** 2, dim=[2, 3], keepdim=True) + self.eps) | |
return x / norm | |
class UpsamplingLayer(nn.Module): | |
"""Implements the upsampling layer. | |
Basically, this layer can be used to upsample feature maps with nearest | |
neighbor interpolation. | |
""" | |
def __init__(self, scale_factor=2): | |
super().__init__() | |
self.scale_factor = scale_factor | |
def forward(self, x): | |
if self.scale_factor <= 1: | |
return x | |
return F.interpolate(x, scale_factor=self.scale_factor, mode='nearest') | |
class Blur(torch.autograd.Function): | |
"""Defines blur operation with customized gradient computation.""" | |
def forward(ctx, x, kernel): | |
ctx.save_for_backward(kernel) | |
y = F.conv2d(input=x, | |
weight=kernel, | |
bias=None, | |
stride=1, | |
padding=1, | |
groups=x.shape[1]) | |
return y | |
def backward(ctx, dy): | |
kernel, = ctx.saved_tensors | |
dx = F.conv2d(input=dy, | |
weight=kernel.flip((2, 3)), | |
bias=None, | |
stride=1, | |
padding=1, | |
groups=dy.shape[1]) | |
return dx, None, None | |
class BlurLayer(nn.Module): | |
"""Implements the blur layer.""" | |
def __init__(self, | |
channels, | |
kernel=(1, 2, 1), | |
normalize=True): | |
super().__init__() | |
kernel = np.array(kernel, dtype=np.float32).reshape(1, -1) | |
kernel = kernel.T.dot(kernel) | |
if normalize: | |
kernel /= np.sum(kernel) | |
kernel = kernel[np.newaxis, np.newaxis] | |
kernel = np.tile(kernel, [channels, 1, 1, 1]) | |
self.register_buffer('kernel', torch.from_numpy(kernel)) | |
def forward(self, x): | |
return Blur.apply(x, self.kernel) | |
class NoiseApplyingLayer(nn.Module): | |
"""Implements the noise applying layer.""" | |
def __init__(self, resolution, channels, noise_type='spatial'): | |
super().__init__() | |
self.noise_type = noise_type.lower() | |
self.res = resolution | |
self.channels = channels | |
if self.noise_type == 'spatial': | |
self.register_buffer('noise', torch.randn(1, 1, self.res, self.res)) | |
self.weight = nn.Parameter(torch.zeros(self.channels)) | |
elif self.noise_type == 'channel': | |
self.register_buffer('noise', torch.randn(1, self.channels, 1, 1)) | |
self.weight = nn.Parameter(torch.zeros(self.res, self.res)) | |
else: | |
raise NotImplementedError(f'Not implemented noise type: ' | |
f'`{self.noise_type}`!') | |
def forward(self, x, randomize_noise=False): | |
if x.ndim != 4: | |
raise ValueError(f'The input tensor should be with shape ' | |
f'[batch_size, channel, height, width], ' | |
f'but `{x.shape}` is received!') | |
if randomize_noise: | |
if self.noise_type == 'spatial': | |
noise = torch.randn(x.shape[0], 1, self.res, self.res).to(x) | |
elif self.noise_type == 'channel': | |
noise = torch.randn(x.shape[0], self.channels, 1, 1).to(x) | |
else: | |
noise = self.noise | |
if self.noise_type == 'spatial': | |
x = x + noise * self.weight.view(1, self.channels, 1, 1) | |
elif self.noise_type == 'channel': | |
x = x + noise * self.weight.view(1, 1, self.res, self.res) | |
return x | |
class StyleModLayer(nn.Module): | |
"""Implements the style modulation layer.""" | |
def __init__(self, | |
w_space_dim, | |
out_channels, | |
use_wscale=True): | |
super().__init__() | |
self.w_space_dim = w_space_dim | |
self.out_channels = out_channels | |
weight_shape = (self.out_channels * 2, self.w_space_dim) | |
wscale = _STYLEMOD_WSCALE_GAIN / np.sqrt(self.w_space_dim) | |
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 | |
self.bias = nn.Parameter(torch.zeros(self.out_channels * 2)) | |
self.space_of_latent = 'w' | |
def forward_style(self, w): | |
"""Gets style code from the given input. | |
More specifically, if the input is from W-Space, it will be projected by | |
an affine transformation. If it is from the Style Space (Y-Space), no | |
operation is required. | |
NOTE: For codes from Y-Space, we use slicing to make sure the dimension | |
is correct, in case that the code is padded before fed into this layer. | |
""" | |
if self.space_of_latent == 'w': | |
if w.ndim != 2 or w.shape[1] != self.w_space_dim: | |
raise ValueError(f'The input tensor should be with shape ' | |
f'[batch_size, w_space_dim], where ' | |
f'`w_space_dim` equals to ' | |
f'{self.w_space_dim}!\n' | |
f'But `{w.shape}` is received!') | |
style = F.linear(w, | |
weight=self.weight * self.wscale, | |
bias=self.bias) | |
elif self.space_of_latent == 'y': | |
if w.ndim != 2 or w.shape[1] < 2 * self.out_channels: | |
raise ValueError(f'The input tensor should be with shape ' | |
f'[batch_size, y_space_dim], where ' | |
f'`y_space_dim` equals to ' | |
f'{2 * self.out_channels}!\n' | |
f'But `{w.shape}` is received!') | |
style = w[:, :2 * self.out_channels] | |
return style | |
def forward(self, x, w): | |
style = self.forward_style(w) | |
style_split = style.view(-1, 2, self.out_channels, 1, 1) | |
x = x * (style_split[:, 0] + 1) + style_split[:, 1] | |
return x, style | |
class ConvBlock(nn.Module): | |
"""Implements the normal convolutional block. | |
Basically, this block executes upsampling layer (if needed), convolutional | |
layer, blurring layer, noise applying layer, activation layer, instance | |
normalization layer, and style modulation layer in sequence. | |
""" | |
def __init__(self, | |
in_channels, | |
out_channels, | |
resolution, | |
w_space_dim, | |
position=None, | |
kernel_size=3, | |
stride=1, | |
padding=1, | |
add_bias=True, | |
upsample=False, | |
fused_scale=False, | |
use_wscale=True, | |
wscale_gain=_WSCALE_GAIN, | |
lr_mul=1.0, | |
activation_type='lrelu', | |
noise_type='spatial'): | |
"""Initializes with block settings. | |
Args: | |
in_channels: Number of channels of the input tensor. | |
out_channels: Number of channels of the output tensor. | |
resolution: Resolution of the output tensor. | |
w_space_dim: Dimension of W space for style modulation. | |
position: Position of the layer. `const_init`, `last` would lead to | |
different behavior. (default: None) | |
kernel_size: Size of the convolutional kernels. (default: 3) | |
stride: Stride parameter for convolution operation. (default: 1) | |
padding: Padding parameter for convolution operation. (default: 1) | |
add_bias: Whether to add bias onto the convolutional result. | |
(default: True) | |
upsample: Whether to upsample the input tensor before convolution. | |
(default: False) | |
fused_scale: Whether to fused `upsample` and `conv2d` together, | |
resulting in `conv2d_transpose`. (default: False) | |
use_wscale: Whether to use weight scaling. (default: True) | |
wscale_gain: Gain factor for weight scaling. (default: _WSCALE_GAIN) | |
lr_mul: Learning multiplier for both weight and bias. (default: 1.0) | |
activation_type: Type of activation. Support `linear` and `lrelu`. | |
(default: `lrelu`) | |
noise_type: Type of noise added to the feature map after the | |
convolution (if needed). Support `spatial` and `channel`. | |
(default: `spatial`) | |
Raises: | |
NotImplementedError: If the `activation_type` is not supported. | |
""" | |
super().__init__() | |
self.position = position | |
if add_bias: | |
self.bias = nn.Parameter(torch.zeros(out_channels)) | |
self.bscale = lr_mul | |
else: | |
self.bias = None | |
if activation_type == 'linear': | |
self.activate = nn.Identity() | |
elif activation_type == 'lrelu': | |
self.activate = nn.LeakyReLU(negative_slope=0.2, inplace=True) | |
else: | |
raise NotImplementedError(f'Not implemented activation function: ' | |
f'`{activation_type}`!') | |
if self.position != 'last': | |
self.apply_noise = NoiseApplyingLayer( | |
resolution, out_channels, noise_type=noise_type) | |
self.normalize = InstanceNormLayer() | |
self.style = StyleModLayer(w_space_dim, out_channels, use_wscale) | |
if self.position == 'const_init': | |
self.const = nn.Parameter( | |
torch.ones(1, in_channels, resolution, resolution)) | |
return | |
self.blur = BlurLayer(out_channels) if upsample else nn.Identity() | |
if upsample and not fused_scale: | |
self.upsample = UpsamplingLayer() | |
else: | |
self.upsample = nn.Identity() | |
if upsample and fused_scale: | |
self.use_conv2d_transpose = True | |
self.stride = 2 | |
self.padding = 1 | |
else: | |
self.use_conv2d_transpose = False | |
self.stride = stride | |
self.padding = padding | |
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) / lr_mul) | |
self.wscale = wscale * lr_mul | |
else: | |
self.weight = nn.Parameter( | |
torch.randn(*weight_shape) * wscale / lr_mul) | |
self.wscale = lr_mul | |
def forward(self, x, w, randomize_noise=False): | |
if self.position != 'const_init': | |
x = self.upsample(x) | |
weight = self.weight * self.wscale | |
if self.use_conv2d_transpose: | |
weight = F.pad(weight, (1, 1, 1, 1, 0, 0, 0, 0), 'constant', 0) | |
weight = (weight[:, :, 1:, 1:] + weight[:, :, :-1, 1:] + | |
weight[:, :, 1:, :-1] + weight[:, :, :-1, :-1]) | |
weight = weight.permute(1, 0, 2, 3) | |
x = F.conv_transpose2d(x, | |
weight=weight, | |
bias=None, | |
stride=self.stride, | |
padding=self.padding) | |
else: | |
x = F.conv2d(x, | |
weight=weight, | |
bias=None, | |
stride=self.stride, | |
padding=self.padding) | |
x = self.blur(x) | |
else: | |
x = self.const.repeat(w.shape[0], 1, 1, 1) | |
bias = self.bias * self.bscale if self.bias is not None else None | |
if self.position == 'last': | |
if bias is not None: | |
x = x + bias.view(1, -1, 1, 1) | |
return x | |
x = self.apply_noise(x, randomize_noise) | |
if bias is not None: | |
x = x + bias.view(1, -1, 1, 1) | |
x = self.activate(x) | |
x = self.normalize(x) | |
x, style = self.style(x, w) | |
return x, style | |
class DenseBlock(nn.Module): | |
"""Implements the dense block. | |
Basically, this block executes fully-connected layer and activation layer. | |
""" | |
def __init__(self, | |
in_channels, | |
out_channels, | |
add_bias=True, | |
use_wscale=True, | |
wscale_gain=_WSCALE_GAIN, | |
lr_mul=1.0, | |
activation_type='lrelu'): | |
"""Initializes with block 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. | |
(default: True) | |
use_wscale: Whether to use weight scaling. (default: True) | |
wscale_gain: Gain factor for weight scaling. (default: _WSCALE_GAIN) | |
lr_mul: Learning multiplier for both weight and bias. (default: 1.0) | |
activation_type: Type of activation. Support `linear` and `lrelu`. | |
(default: `lrelu`) | |
Raises: | |
NotImplementedError: If the `activation_type` is not supported. | |
""" | |
super().__init__() | |
weight_shape = (out_channels, in_channels) | |
wscale = wscale_gain / np.sqrt(in_channels) | |
if use_wscale: | |
self.weight = nn.Parameter(torch.randn(*weight_shape) / lr_mul) | |
self.wscale = wscale * lr_mul | |
else: | |
self.weight = nn.Parameter( | |
torch.randn(*weight_shape) * wscale / lr_mul) | |
self.wscale = lr_mul | |
if add_bias: | |
self.bias = nn.Parameter(torch.zeros(out_channels)) | |
self.bscale = lr_mul | |
else: | |
self.bias = None | |
if activation_type == 'linear': | |
self.activate = nn.Identity() | |
elif activation_type == 'lrelu': | |
self.activate = nn.LeakyReLU(negative_slope=0.2, inplace=True) | |
else: | |
raise NotImplementedError(f'Not implemented activation function: ' | |
f'`{activation_type}`!') | |
def forward(self, x): | |
if x.ndim != 2: | |
x = x.view(x.shape[0], -1) | |
bias = self.bias * self.bscale if self.bias is not None else None | |
x = F.linear(x, weight=self.weight * self.wscale, bias=bias) | |
x = self.activate(x) | |
return x | |