From 81abb9141c893b779eae7508139043ecbd9ed149 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Thu, 16 Jun 2022 15:28:52 +0000 Subject: [PATCH] add ae sub folder --- stylegan2_ae/__init__.py | 0 stylegan2_ae/discriminator.py | 152 +++++++++ stylegan2_ae/generator.py | 113 +++++++ stylegan2_ae/layers/__init__.py | 0 stylegan2_ae/layers/bias_act.py | 73 +++++ stylegan2_ae/layers/commons.py | 20 ++ stylegan2_ae/layers/conv.py | 59 ++++ stylegan2_ae/layers/cuda/__init__.py | 7 + stylegan2_ae/layers/cuda/custom_ops.py | 169 ++++++++++ stylegan2_ae/layers/cuda/fused_bias_act.cu | 188 +++++++++++ stylegan2_ae/layers/cuda/fused_bias_act.py | 196 +++++++++++ stylegan2_ae/layers/cuda/upfirdn_2d.cu | 326 ++++++++++++++++++ stylegan2_ae/layers/cuda/upfirdn_2d.py | 364 +++++++++++++++++++++ stylegan2_ae/layers/cuda/upfirdn_2d_v2.py | 146 +++++++++ stylegan2_ae/layers/dense.py | 37 +++ stylegan2_ae/layers/embedding.py | 135 ++++++++ stylegan2_ae/layers/from_rgb.py | 28 ++ stylegan2_ae/layers/mapping_block.py | 48 +++ stylegan2_ae/layers/mini_batch_std.py | 33 ++ stylegan2_ae/layers/modulated_conv2d.py | 105 ++++++ stylegan2_ae/layers/noise.py | 25 ++ stylegan2_ae/layers/synthesis_block.py | 162 +++++++++ stylegan2_ae/layers/to_rgb.py | 30 ++ stylegan2_ae/utils.py | 112 +++++++ 24 files changed, 2528 insertions(+) create mode 100644 stylegan2_ae/__init__.py create mode 100644 stylegan2_ae/discriminator.py create mode 100644 stylegan2_ae/generator.py create mode 100644 stylegan2_ae/layers/__init__.py create mode 100644 stylegan2_ae/layers/bias_act.py create mode 100644 stylegan2_ae/layers/commons.py create mode 100644 stylegan2_ae/layers/conv.py create mode 100644 stylegan2_ae/layers/cuda/__init__.py create mode 100644 stylegan2_ae/layers/cuda/custom_ops.py create mode 100644 stylegan2_ae/layers/cuda/fused_bias_act.cu create mode 100644 stylegan2_ae/layers/cuda/fused_bias_act.py create mode 100644 stylegan2_ae/layers/cuda/upfirdn_2d.cu create mode 100644 stylegan2_ae/layers/cuda/upfirdn_2d.py create mode 100644 stylegan2_ae/layers/cuda/upfirdn_2d_v2.py create mode 100644 stylegan2_ae/layers/dense.py create mode 100644 stylegan2_ae/layers/embedding.py create mode 100644 stylegan2_ae/layers/from_rgb.py create mode 100644 stylegan2_ae/layers/mapping_block.py create mode 100644 stylegan2_ae/layers/mini_batch_std.py create mode 100644 stylegan2_ae/layers/modulated_conv2d.py create mode 100644 stylegan2_ae/layers/noise.py create mode 100644 stylegan2_ae/layers/synthesis_block.py create mode 100644 stylegan2_ae/layers/to_rgb.py create mode 100644 stylegan2_ae/utils.py diff --git a/stylegan2_ae/__init__.py b/stylegan2_ae/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/stylegan2_ae/discriminator.py b/stylegan2_ae/discriminator.py new file mode 100644 index 0000000..4544d2c --- /dev/null +++ b/stylegan2_ae/discriminator.py @@ -0,0 +1,152 @@ +import tensorflow as tf + +from stylegan2.layers.dense import Dense +from stylegan2.layers.conv import Conv2D +from stylegan2.layers.bias_act import BiasAct +from stylegan2.layers.from_rgb import FromRGB +from stylegan2.layers.mini_batch_std import MinibatchStd + + +class DiscriminatorBlock(tf.keras.layers.Layer): + def __init__(self, n_f0, n_f1, res, **kwargs): + super(DiscriminatorBlock, self).__init__(**kwargs) + self.gain = 1.0 + self.lrmul = 1.0 + self.n_f0 = n_f0 + self.n_f1 = n_f1 + self.res = res + self.resnet_scale = 1. / tf.sqrt(2.) + + # conv_0 + self.conv_0 = Conv2D(in_res=res, in_fmaps=self.n_f0, fmaps=self.n_f0, kernel=3, up=False, down=False, + resample_kernel=None, gain=self.gain, lrmul=self.lrmul, name='conv_0') + self.apply_bias_act_0 = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_0') + + # conv_1 down + self.conv_1 = Conv2D(in_res=res, in_fmaps=self.n_f0, fmaps=self.n_f1, kernel=3, up=False, down=True, + resample_kernel=[1, 3, 3, 1], gain=self.gain, lrmul=self.lrmul, name='conv_1') + self.apply_bias_act_1 = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_1') + + # resnet skip + self.conv_skip = Conv2D(in_res=res, in_fmaps=self.n_f0, fmaps=self.n_f1, kernel=1, up=False, down=True, + resample_kernel=[1, 3, 3, 1], gain=self.gain, lrmul=self.lrmul, name='skip') + + def call(self, inputs, training=None, mask=None): + x = inputs + residual = x + + # conv0 + x = self.conv_0(x) + x = self.apply_bias_act_0(x) + + # conv1 down + x = self.conv_1(x) + x = self.apply_bias_act_1(x) + + # resnet skip + residual = self.conv_skip(residual) + x = (x + residual) * self.resnet_scale + return x + + def get_config(self): + config = super(DiscriminatorBlock, self).get_config() + config.update({ + 'n_f0': self.n_f0, + 'n_f1': self.n_f1, + 'gain': self.gain, + 'lrmul': self.lrmul, + 'res': self.res, + 'resnet_scale': self.resnet_scale, + }) + return config + + +class DiscriminatorLastBlock(tf.keras.layers.Layer): + def __init__(self, n_f0, n_f1, res, **kwargs): + super(DiscriminatorLastBlock, self).__init__(**kwargs) + self.gain = 1.0 + self.lrmul = 1.0 + self.n_f0 = n_f0 + self.n_f1 = n_f1 + self.res = res + + self.minibatch_std = MinibatchStd(group_size=4, num_new_features=1, name='minibatchstd') + + # conv_0 + self.conv_0 = Conv2D(in_res=res, in_fmaps=self.n_f0 + 1, fmaps=self.n_f0, kernel=3, up=False, down=False, + resample_kernel=None, gain=self.gain, lrmul=self.lrmul, name='conv_0') + self.apply_bias_act_0 = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_0') + + # dense_1 + self.dense_1 = Dense(self.n_f1, gain=self.gain, lrmul=self.lrmul, name='dense_1') + self.apply_bias_act_1 = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_1') + + def call(self, x, training=None, mask=None): + x = self.minibatch_std(x) + + # conv_0 + x = self.conv_0(x) + x = self.apply_bias_act_0(x) + + # dense_1 + x = self.dense_1(x) + x = self.apply_bias_act_1(x) + return x + + def get_config(self): + config = super(DiscriminatorLastBlock, self).get_config() + config.update({ + 'n_f0': self.n_f0, + 'n_f1': self.n_f1, + 'gain': self.gain, + 'lrmul': self.lrmul, + 'res': self.res, + }) + return config + + +class Discriminator(tf.keras.Model): + def __init__(self, d_params, **kwargs): + super(Discriminator, self).__init__(**kwargs) + # discriminator's (resolutions and featuremaps) are reversed against generator's + self.labels_dim = d_params['labels_dim'] + self.r_resolutions = d_params['resolutions'][::-1] + self.r_featuremaps = d_params['featuremaps'][::-1] + + # stack discriminator blocks + res0, n_f0 = self.r_resolutions[0], self.r_featuremaps[0] + self.initial_fromrgb = FromRGB(fmaps=n_f0, res=res0, name='{:d}x{:d}/FromRGB'.format(res0, res0)) + self.blocks = list() + for index, (res0, n_f0) in enumerate(zip(self.r_resolutions[:-1], self.r_featuremaps[:-1])): + n_f1 = self.r_featuremaps[index + 1] + self.blocks.append(DiscriminatorBlock(n_f0=n_f0, n_f1=n_f1, res=res0, name='{:d}x{:d}'.format(res0, res0))) + + # set last discriminator block + res = self.r_resolutions[-1] + n_f0, n_f1 = self.r_featuremaps[-2], self.r_featuremaps[-1] + self.last_block = DiscriminatorLastBlock(n_f0, n_f1, res, name='{:d}x{:d}'.format(res, res)) + + # set last dense layer + self.last_dense = Dense(max(self.labels_dim, 1), gain=1.0, lrmul=1.0, name='last_dense') + self.last_bias = BiasAct(lrmul=1.0, act='linear', name='last_bias') + + + def call(self, inputs, training=None, mask=None): + images = inputs + + x = self.initial_fromrgb(images) + for block in self.blocks: + x = block(x) + + features = x + + # score output + x = self.last_block(features) + x = self.last_dense(x) + x = self.last_bias(x) + + scores_out = x + return scores_out + + def compute_output_shape(self, input_shape): + return input_shape[0][0], max(self.labels_dim, 1) diff --git a/stylegan2_ae/generator.py b/stylegan2_ae/generator.py new file mode 100644 index 0000000..7159d40 --- /dev/null +++ b/stylegan2_ae/generator.py @@ -0,0 +1,113 @@ +import tensorflow as tf + +from stylegan2_ae.utils import lerp +from stylegan2_ae.layers.mapping_block import Mapping +from stylegan2_ae.layers.synthesis_block import Synthesis + + +class Generator(tf.keras.Model): + def __init__(self, g_params, **kwargs): + super(Generator, self).__init__(**kwargs) + + self.z_dim = g_params['z_dim'] + self.w_dim = g_params['w_dim'] + self.labels_dim = g_params['labels_dim'] + self.n_mapping = g_params['n_mapping'] + self.resolutions = g_params['resolutions'] + self.featuremaps = g_params['featuremaps'] + self.w_ema_decay = 0.995 + self.style_mixing_prob = 0.9 + + self.n_broadcast = len(self.resolutions) * 2 + self.mixing_layer_indices = tf.range(self.n_broadcast, dtype=tf.int32)[tf.newaxis, :, tf.newaxis] + + self.g_mapping = Mapping(self.w_dim, self.labels_dim, self.n_mapping, name='g_mapping') + self.broadcast = tf.keras.layers.Lambda(lambda x: tf.tile(x[:, tf.newaxis], [1, self.n_broadcast, 1])) + self.synthesis = Synthesis(self.resolutions, self.featuremaps, name='g_synthesis') + + def build(self, input_shape): + # w_avg + self.w_avg = tf.Variable(tf.zeros(shape=[self.w_dim], dtype=tf.dtypes.float32), name='w_avg', trainable=False, + synchronization=tf.VariableSynchronization.ON_READ, + aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA) + + @tf.function + def set_as_moving_average_of(self, src_net): + beta, beta_nontrainable = 0.99, 0.0 + + for cw, sw in zip(self.weights, src_net.weights): + assert sw.shape == cw.shape + # print('{} <=> {}'.format(cw.name, sw.name)) + + if 'w_avg' in cw.name: + cw.assign(lerp(sw, cw, beta_nontrainable)) + else: + cw.assign(lerp(sw, cw, beta)) + return + + def update_moving_average_of_w(self, w_broadcasted): + # compute average of current w + batch_avg = tf.reduce_mean(w_broadcasted[:, 0], axis=0) + + # compute moving average of w and update(assign) w_avg + self.w_avg.assign(lerp(batch_avg, self.w_avg, self.w_ema_decay)) + return + + def style_mixing_regularization(self, latents1, w_broadcasted1): + # get another w and broadcast it + latents2 = tf.random.normal(shape=tf.shape(latents1), dtype=tf.dtypes.float32) + dlatents2 = self.g_mapping(latents2) + w_broadcasted2 = self.broadcast(dlatents2) + + # find mixing limit index + # mixing_cutoff_index = tf.cond( + # pred=tf.less(tf.random.uniform([], 0.0, 1.0), self.style_mixing_prob), + # true_fn=lambda: tf.random.uniform([], 1, self.n_broadcast, dtype=tf.dtypes.int32), + # false_fn=lambda: tf.constant(self.n_broadcast, dtype=tf.dtypes.int32)) + if tf.random.uniform([], 0.0, 1.0) < self.style_mixing_prob: + mixing_cutoff_index = tf.random.uniform([], 1, self.n_broadcast, dtype=tf.dtypes.int32) + else: + mixing_cutoff_index = tf.constant(self.n_broadcast, dtype=tf.dtypes.int32) + + # mix it + mixed_w_broadcasted = tf.where( + condition=tf.broadcast_to(self.mixing_layer_indices < mixing_cutoff_index, tf.shape(w_broadcasted1)), + x=w_broadcasted1, + y=w_broadcasted2) + return mixed_w_broadcasted + + def truncation_trick(self, w_broadcasted, truncation_psi, truncation_cutoff=None): + ones = tf.ones_like(self.mixing_layer_indices, dtype=tf.float32) + tpsi = ones * truncation_psi + if truncation_cutoff is None: + truncation_coefs = tpsi + else: + indices = tf.range(self.n_broadcast) + truncation_coefs = tf.where(condition=tf.less(indices, truncation_cutoff), x=tpsi, y=ones) + + truncated_w_broadcasted = lerp(self.w_avg, w_broadcasted, truncation_coefs) + return truncated_w_broadcasted + + def call(self, inputs, ret_w_broadcasted=False, truncation_psi=1.0, truncation_cutoff=None, training=None, mask=None): + latents = inputs + + dlatents = self.g_mapping(latents) + w_broadcasted = self.broadcast(dlatents) + + if training: + self.update_moving_average_of_w(w_broadcasted) + w_broadcasted = self.style_mixing_regularization(latents, w_broadcasted) + + if not training: + w_broadcasted = self.truncation_trick(w_broadcasted, truncation_psi, truncation_cutoff) + + image_out = self.synthesis(w_broadcasted) + image_out = image_out + + if ret_w_broadcasted: + return image_out, w_broadcasted + else: + return image_out + + def compute_output_shape(self, input_shape): + return input_shape[0][0], 3, self.resolutions[-1], self.resolutions[-1] diff --git a/stylegan2_ae/layers/__init__.py b/stylegan2_ae/layers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/stylegan2_ae/layers/bias_act.py b/stylegan2_ae/layers/bias_act.py new file mode 100644 index 0000000..bea2284 --- /dev/null +++ b/stylegan2_ae/layers/bias_act.py @@ -0,0 +1,73 @@ +import numpy as np +import tensorflow as tf + + +# the custom cuda compiled fused_bias_act is not compatible with tf.GradientTape() +# I think it's because of 2nd order gradient computation +class BiasAct(tf.keras.layers.Layer): + def __init__(self, lrmul, act, **kwargs): + super(BiasAct, self).__init__(**kwargs) + assert act in ['linear', 'lrelu'] + self.lrmul = lrmul + + if act == 'linear': + self.act = tf.keras.layers.Lambda(lambda x: tf.identity(x)) + self.gain = 1.0 + else: + self.act = tf.keras.layers.LeakyReLU(alpha=0.2) + self.gain = np.sqrt(2) + + def build(self, input_shape): + self.len2 = True if len(input_shape) == 2 else False + b_init = tf.zeros(shape=(input_shape[1],), dtype=tf.dtypes.float32) + self.b = tf.Variable(b_init, name='b', trainable=True) + + def call(self, inputs, training=None, mask=None): + b = self.lrmul * self.b + + if self.len2: + x = inputs + b + else: + x = inputs + tf.reshape(b, shape=[1, -1, 1, 1]) + x = self.act(x) + x = self.gain * x + return x + + def get_config(self): + config = super(BiasAct, self).get_config() + config.update({ + 'lrmul': self.lrmul, + 'gain': self.gain, + 'len2': self.len2, + }) + return config + + +# import tensorflow as tf +# +# +# from stylegan2.layers.cuda.fused_bias_act import fused_bias_act +# +# +# class BiasAct(tf.keras.layers.Layer): +# def __init__(self, lrmul, act, **kwargs): +# super(BiasAct, self).__init__(**kwargs) +# self.lrmul = lrmul +# self.act = act +# +# def build(self, input_shape): +# b_init = tf.zeros(shape=(input_shape[1],), dtype=tf.dtypes.float32) +# self.b = tf.Variable(b_init, name='b', trainable=True) +# +# def call(self, inputs, training=None, mask=None): +# b = self.lrmul * self.b +# x = fused_bias_act(inputs, b=b, act=self.act, alpha=None, gain=None) +# return x +# +# def get_config(self): +# config = super(BiasAct, self).get_config() +# config.update({ +# 'lrmul': self.lrmul, +# 'act': self.act, +# }) +# return config diff --git a/stylegan2_ae/layers/commons.py b/stylegan2_ae/layers/commons.py new file mode 100644 index 0000000..dc90dec --- /dev/null +++ b/stylegan2_ae/layers/commons.py @@ -0,0 +1,20 @@ +import tensorflow as tf + + +def compute_runtime_coef(weight_shape, gain, lrmul): + fan_in = tf.reduce_prod(weight_shape[:-1]) # [kernel, kernel, fmaps_in, fmaps_out] or [in, out] + fan_in = tf.cast(fan_in, dtype=tf.float32) + he_std = gain / tf.sqrt(fan_in) + init_std = 1.0 / lrmul + runtime_coef = he_std * lrmul + return init_std, runtime_coef + + +def lerp(a, b, t): + out = a + (b - a) * t + return out + + +def lerp_clip(a, b, t): + out = a + (b - a) * tf.clip_by_value(t, 0.0, 1.0) + return out diff --git a/stylegan2_ae/layers/conv.py b/stylegan2_ae/layers/conv.py new file mode 100644 index 0000000..7779f6f --- /dev/null +++ b/stylegan2_ae/layers/conv.py @@ -0,0 +1,59 @@ +import tensorflow as tf + +from stylegan2_ae.layers.commons import compute_runtime_coef +from stylegan2_ae.layers.cuda.upfirdn_2d_v2 import upsample_conv_2d, conv_downsample_2d, compute_paddings + + +class Conv2D(tf.keras.layers.Layer): + def __init__(self, in_res, in_fmaps, fmaps, kernel, up, down, resample_kernel, gain, lrmul, **kwargs): + super(Conv2D, self).__init__(**kwargs) + self.in_res = in_res + self.in_fmaps = in_fmaps + self.fmaps = fmaps + self.kernel = kernel + self.gain = gain + self.lrmul = lrmul + self.up = up + self.down = down + # self.resample_kernel = resample_kernel + + self.k, self.pad0, self.pad1 = compute_paddings(resample_kernel, self.kernel, up, down, is_conv=True) + + def build(self, input_shape): + weight_shape = [self.kernel, self.kernel, input_shape[1], self.fmaps] + init_std, self.runtime_coef = compute_runtime_coef(weight_shape, self.gain, self.lrmul) + + # [kernel, kernel, fmaps_in, fmaps_out] + w_init = tf.random.normal(shape=weight_shape, mean=0.0, stddev=init_std) + self.w = tf.Variable(w_init, name='w', trainable=True) + + def call(self, inputs, training=None, mask=None): + x = inputs + w = self.runtime_coef * self.w + + # actual conv + if self.up: + x = upsample_conv_2d(x, self.in_res, w, self.kernel, self.kernel, self.pad0, self.pad1, self.k) + elif self.down: + x = conv_downsample_2d(x, self.in_res, w, self.kernel, self.kernel, self.pad0, self.pad1, self.k) + else: + x = tf.nn.conv2d(x, w, data_format='NCHW', strides=[1, 1, 1, 1], padding='SAME') + return x + + def get_config(self): + config = super(Conv2D, self).get_config() + config.update({ + 'in_res': self.in_res, + 'in_fmaps': self.in_fmaps, + 'fmaps': self.fmaps, + 'kernel': self.kernel, + 'gain': self.gain, + 'lrmul': self.lrmul, + 'up': self.up, + 'down': self.down, + 'k': self.k, + 'pad0': self.pad0, + 'pad1': self.pad1, + 'runtime_coef': self.runtime_coef, + }) + return config diff --git a/stylegan2_ae/layers/cuda/__init__.py b/stylegan2_ae/layers/cuda/__init__.py new file mode 100644 index 0000000..9ab9908 --- /dev/null +++ b/stylegan2_ae/layers/cuda/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2019, NVIDIA Corporation. All rights reserved. +# +# This work is made available under the Nvidia Source Code License-NC. +# To view a copy of this license, visit +# https://nvlabs.github.io/stylegan2/license.html + +# empty diff --git a/stylegan2_ae/layers/cuda/custom_ops.py b/stylegan2_ae/layers/cuda/custom_ops.py new file mode 100644 index 0000000..a55a445 --- /dev/null +++ b/stylegan2_ae/layers/cuda/custom_ops.py @@ -0,0 +1,169 @@ +# Copyright (c) 2019, NVIDIA Corporation. All rights reserved. +# +# This work is made available under the Nvidia Source Code License-NC. +# To view a copy of this license, visit +# https://nvlabs.github.io/stylegan2/license.html + +"""TensorFlow custom ops builder. +""" + +import os +import re +import uuid +import hashlib +import tempfile +import shutil +import tensorflow as tf +from tensorflow.python.client import device_lib # pylint: disable=no-name-in-module + +#---------------------------------------------------------------------------- +# Global options. + +cuda_cache_path = os.path.join(os.path.dirname(__file__), '_cudacache') +cuda_cache_version_tag = 'v1' +do_not_hash_included_headers = False # Speed up compilation by assuming that headers included by the CUDA code never change. Unsafe! +verbose = True # Print status messages to stdout. + +compiler_bindir_search_path = [ + 'C:/Program Files (x86)/Microsoft Visual Studio/2017/Community/VC/Tools/MSVC/14.14.26428/bin/Hostx64/x64', + 'C:/Program Files (x86)/Microsoft Visual Studio/2019/Community/VC/Tools/MSVC/14.23.28105/bin/Hostx64/x64', + 'C:/Program Files (x86)/Microsoft Visual Studio 14.0/vc/bin', +] + +#---------------------------------------------------------------------------- +# Internal helper funcs. + +def _find_compiler_bindir(): + for compiler_path in compiler_bindir_search_path: + if os.path.isdir(compiler_path): + return compiler_path + return None + +def _get_compute_cap(device): + caps_str = device.physical_device_desc + m = re.search('compute capability: (\\d+).(\\d+)', caps_str) + major = m.group(1) + minor = m.group(2) + return (major, minor) + +def _get_cuda_gpu_arch_string(): + gpus = [x for x in device_lib.list_local_devices() if x.device_type == 'GPU'] + if len(gpus) == 0: + raise RuntimeError('No GPU devices found') + (major, minor) = _get_compute_cap(gpus[0]) + return 'sm_%s%s' % (major, minor) + +def _run_cmd(cmd): + with os.popen(cmd) as pipe: + output = pipe.read() + status = pipe.close() + if status is not None: + raise RuntimeError('NVCC returned an error. See below for full command line and output log:\n\n%s\n\n%s' % (cmd, output)) + +def _prepare_nvcc_cli(opts): + cmd = 'nvcc --std=c++14 -DNDEBUG ' + opts.strip() + cmd += ' --disable-warnings' + cmd += ' --include-path "%s"' % tf.sysconfig.get_include() + cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'protobuf_archive', 'src') + cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'com_google_absl') + cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'eigen_archive') + + compiler_bindir = _find_compiler_bindir() + if compiler_bindir is None: + # Require that _find_compiler_bindir succeeds on Windows. Allow + # nvcc to use whatever is the default on Linux. + if os.name == 'nt': + raise RuntimeError('Could not find MSVC/GCC/CLANG installation on this computer. Check compiler_bindir_search_path list in "%s".' % __file__) + else: + cmd += ' --compiler-bindir "%s"' % compiler_bindir + cmd += ' 2>&1' + return cmd + +#---------------------------------------------------------------------------- +# Main entry point. + +_plugin_cache = dict() + +def get_plugin(cuda_file): + cuda_file_base = os.path.basename(cuda_file) + cuda_file_name, cuda_file_ext = os.path.splitext(cuda_file_base) + + # Already in cache? + if cuda_file in _plugin_cache: + return _plugin_cache[cuda_file] + + # Setup plugin. + if verbose: + print('Setting up TensorFlow plugin "%s": ' % cuda_file_base, end='', flush=True) + try: + # Hash CUDA source. + md5 = hashlib.md5() + with open(cuda_file, 'rb') as f: + md5.update(f.read()) + md5.update(b'\n') + + # Hash headers included by the CUDA code by running it through the preprocessor. + if not do_not_hash_included_headers: + if verbose: + print('Preprocessing... ', end='', flush=True) + with tempfile.TemporaryDirectory() as tmp_dir: + tmp_file = os.path.join(tmp_dir, cuda_file_name + '_tmp' + cuda_file_ext) + _run_cmd(_prepare_nvcc_cli('"%s" --preprocess -o "%s" --keep --keep-dir "%s"' % (cuda_file, tmp_file, tmp_dir))) + with open(tmp_file, 'rb') as f: + bad_file_str = ('"' + cuda_file.replace('\\', '/') + '"').encode('utf-8') # __FILE__ in error check macros + good_file_str = ('"' + cuda_file_base + '"').encode('utf-8') + for ln in f: + if not ln.startswith(b'# ') and not ln.startswith(b'#line '): # ignore line number pragmas + ln = ln.replace(bad_file_str, good_file_str) + md5.update(ln) + md5.update(b'\n') + + # Select compiler options. + compile_opts = '' + if os.name == 'nt': + compile_opts += '"%s"' % os.path.join(tf.sysconfig.get_lib(), 'python', '_pywrap_tensorflow_internal.lib') + elif os.name == 'posix': + compile_opts += '"%s"' % os.path.join(tf.sysconfig.get_lib(), 'python', '_pywrap_tensorflow_internal.so') + compile_opts += ' --compiler-options \'-fPIC -D_GLIBCXX_USE_CXX11_ABI=1\'' + else: + assert False # not Windows or Linux, w00t? + compile_opts += ' --gpu-architecture=%s' % _get_cuda_gpu_arch_string() + compile_opts += ' --use_fast_math' + nvcc_cmd = _prepare_nvcc_cli(compile_opts) + + # Hash build configuration. + md5.update(('nvcc_cmd: ' + nvcc_cmd).encode('utf-8') + b'\n') + md5.update(('tf.VERSION: ' + tf.version.VERSION).encode('utf-8') + b'\n') + md5.update(('cuda_cache_version_tag: ' + cuda_cache_version_tag).encode('utf-8') + b'\n') + + # Compile if not already compiled. + bin_file_ext = '.dll' if os.name == 'nt' else '.so' + bin_file = os.path.join(cuda_cache_path, cuda_file_name + '_' + md5.hexdigest() + bin_file_ext) + if not os.path.isfile(bin_file): + if verbose: + print('Compiling... ', end='', flush=True) + with tempfile.TemporaryDirectory() as tmp_dir: + tmp_file = os.path.join(tmp_dir, cuda_file_name + '_tmp' + bin_file_ext) + _run_cmd(nvcc_cmd + ' "%s" --shared -o "%s" --keep --keep-dir "%s"' % (cuda_file, tmp_file, tmp_dir)) + os.makedirs(cuda_cache_path, exist_ok=True) + intermediate_file = os.path.join(cuda_cache_path, cuda_file_name + '_' + uuid.uuid4().hex + '_tmp' + bin_file_ext) + shutil.copyfile(tmp_file, intermediate_file) + os.rename(intermediate_file, bin_file) # atomic + + # Load. + if verbose: + print('Loading... ', end='', flush=True) + plugin = tf.load_op_library(bin_file) + + # Add to cache. + _plugin_cache[cuda_file] = plugin + if verbose: + print('Done.', flush=True) + return plugin + + except: + if verbose: + print('Failed!', flush=True) + raise + +#---------------------------------------------------------------------------- diff --git a/stylegan2_ae/layers/cuda/fused_bias_act.cu b/stylegan2_ae/layers/cuda/fused_bias_act.cu new file mode 100644 index 0000000..1102f62 --- /dev/null +++ b/stylegan2_ae/layers/cuda/fused_bias_act.cu @@ -0,0 +1,188 @@ +// Copyright (c) 2019, NVIDIA Corporation. All rights reserved. +// +// This work is made available under the Nvidia Source Code License-NC. +// To view a copy of this license, visit +// https://nvlabs.github.io/stylegan2/license.html + +#define EIGEN_USE_GPU +#define __CUDA_INCLUDE_COMPILER_INTERNAL_HEADERS__ +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/shape_inference.h" +#include + +using namespace tensorflow; +using namespace tensorflow::shape_inference; + +#define OP_CHECK_CUDA_ERROR(CTX, CUDA_CALL) do { cudaError_t err = CUDA_CALL; OP_REQUIRES(CTX, err == cudaSuccess, errors::Internal(cudaGetErrorName(err))); } while (false) + +//------------------------------------------------------------------------ +// CUDA kernel. + +template +struct FusedBiasActKernelParams +{ + const T* x; // [sizeX] + const T* b; // [sizeB] or NULL + const T* ref; // [sizeX] or NULL + T* y; // [sizeX] + + int grad; + int axis; + int act; + float alpha; + float gain; + + int sizeX; + int sizeB; + int stepB; + int loopX; +}; + +template +static __global__ void FusedBiasActKernel(const FusedBiasActKernelParams p) +{ + const float expRange = 80.0f; + const float halfExpRange = 40.0f; + const float seluScale = 1.0507009873554804934193349852946f; + const float seluAlpha = 1.6732632423543772848170429916717f; + + // Loop over elements. + int xi = blockIdx.x * p.loopX * blockDim.x + threadIdx.x; + for (int loopIdx = 0; loopIdx < p.loopX && xi < p.sizeX; loopIdx++, xi += blockDim.x) + { + // Load and apply bias. + float x = (float)p.x[xi]; + if (p.b) + x += (float)p.b[(xi / p.stepB) % p.sizeB]; + float ref = (p.ref) ? (float)p.ref[xi] : 0.0f; + if (p.gain != 0.0f & p.act != 9) + ref /= p.gain; + + // Evaluate activation func. + float y; + switch (p.act * 10 + p.grad) + { + // linear + default: + case 10: y = x; break; + case 11: y = x; break; + case 12: y = 0.0f; break; + + // relu + case 20: y = (x > 0.0f) ? x : 0.0f; break; + case 21: y = (ref > 0.0f) ? x : 0.0f; break; + case 22: y = 0.0f; break; + + // lrelu + case 30: y = (x > 0.0f) ? x : x * p.alpha; break; + case 31: y = (ref > 0.0f) ? x : x * p.alpha; break; + case 32: y = 0.0f; break; + + // tanh + case 40: { float c = expf(x); float d = 1.0f / c; y = (x < -expRange) ? -1.0f : (x > expRange) ? 1.0f : (c - d) / (c + d); } break; + case 41: y = x * (1.0f - ref * ref); break; + case 42: y = x * (1.0f - ref * ref) * (-2.0f * ref); break; + + // sigmoid + case 50: y = (x < -expRange) ? 0.0f : 1.0f / (expf(-x) + 1.0f); break; + case 51: y = x * ref * (1.0f - ref); break; + case 52: y = x * ref * (1.0f - ref) * (1.0f - 2.0f * ref); break; + + // elu + case 60: y = (x >= 0.0f) ? x : expf(x) - 1.0f; break; + case 61: y = (ref >= 0.0f) ? x : x * (ref + 1.0f); break; + case 62: y = (ref >= 0.0f) ? 0.0f : x * (ref + 1.0f); break; + + // selu + case 70: y = (x >= 0.0f) ? seluScale * x : (seluScale * seluAlpha) * (expf(x) - 1.0f); break; + case 71: y = (ref >= 0.0f) ? x * seluScale : x * (ref + seluScale * seluAlpha); break; + case 72: y = (ref >= 0.0f) ? 0.0f : x * (ref + seluScale * seluAlpha); break; + + // softplus + case 80: y = (x > expRange) ? x : logf(expf(x) + 1.0f); break; + case 81: y = x * (1.0f - expf(-ref)); break; + case 82: { float c = expf(-ref); y = x * c * (1.0f - c); } break; + + // swish + case 90: y = (x < -expRange) ? 0.0f : x / (expf(-x) + 1.0f); break; + case 91: { float c = expf(ref); float d = c + 1.0f; y = (ref > halfExpRange) ? x : x * c * (ref + d) / (d * d); } break; + case 92: { float c = expf(ref); float d = c + 1.0f; y = (ref > halfExpRange) ? 0.0f : x * c * (ref * (2.0f - d) + 2.0f * d) / (d * d * d); } break; + } + + // Apply gain and store. + p.y[xi] = (T)(y * p.gain); + } +} + +//------------------------------------------------------------------------ +// TensorFlow op. + +template +struct FusedBiasActOp : public OpKernel +{ + FusedBiasActKernelParams m_attribs; + + FusedBiasActOp(OpKernelConstruction* ctx) : OpKernel(ctx) + { + memset(&m_attribs, 0, sizeof(m_attribs)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("grad", &m_attribs.grad)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("axis", &m_attribs.axis)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("act", &m_attribs.act)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("alpha", &m_attribs.alpha)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("gain", &m_attribs.gain)); + OP_REQUIRES(ctx, m_attribs.grad >= 0, errors::InvalidArgument("grad must be non-negative")); + OP_REQUIRES(ctx, m_attribs.axis >= 0, errors::InvalidArgument("axis must be non-negative")); + OP_REQUIRES(ctx, m_attribs.act >= 0, errors::InvalidArgument("act must be non-negative")); + } + + void Compute(OpKernelContext* ctx) + { + FusedBiasActKernelParams p = m_attribs; + cudaStream_t stream = ctx->eigen_device().stream(); + + const Tensor& x = ctx->input(0); // [...] + const Tensor& b = ctx->input(1); // [sizeB] or [0] + const Tensor& ref = ctx->input(2); // x.shape or [0] + p.x = x.flat().data(); + p.b = (b.NumElements()) ? b.flat().data() : NULL; + p.ref = (ref.NumElements()) ? ref.flat().data() : NULL; + OP_REQUIRES(ctx, b.NumElements() == 0 || m_attribs.axis < x.dims(), errors::InvalidArgument("axis out of bounds")); + OP_REQUIRES(ctx, b.dims() == 1, errors::InvalidArgument("b must have rank 1")); + OP_REQUIRES(ctx, b.NumElements() == 0 || b.NumElements() == x.dim_size(m_attribs.axis), errors::InvalidArgument("b has wrong number of elements")); + OP_REQUIRES(ctx, ref.NumElements() == ((p.grad == 0) ? 0 : x.NumElements()), errors::InvalidArgument("ref has wrong number of elements")); + OP_REQUIRES(ctx, x.NumElements() <= kint32max, errors::InvalidArgument("x is too large")); + + p.sizeX = (int)x.NumElements(); + p.sizeB = (int)b.NumElements(); + p.stepB = 1; + for (int i = m_attribs.axis + 1; i < x.dims(); i++) + p.stepB *= (int)x.dim_size(i); + + Tensor* y = NULL; // x.shape + OP_REQUIRES_OK(ctx, ctx->allocate_output(0, x.shape(), &y)); + p.y = y->flat().data(); + + p.loopX = 4; + int blockSize = 4 * 32; + int gridSize = (p.sizeX - 1) / (p.loopX * blockSize) + 1; + void* args[] = {&p}; + OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel((void*)FusedBiasActKernel, gridSize, blockSize, args, 0, stream)); + } +}; + +REGISTER_OP("FusedBiasAct") + .Input ("x: T") + .Input ("b: T") + .Input ("ref: T") + .Output ("y: T") + .Attr ("T: {float, half}") + .Attr ("grad: int = 0") + .Attr ("axis: int = 1") + .Attr ("act: int = 0") + .Attr ("alpha: float = 0.0") + .Attr ("gain: float = 1.0"); +REGISTER_KERNEL_BUILDER(Name("FusedBiasAct").Device(DEVICE_GPU).TypeConstraint("T"), FusedBiasActOp); +REGISTER_KERNEL_BUILDER(Name("FusedBiasAct").Device(DEVICE_GPU).TypeConstraint("T"), FusedBiasActOp); + +//------------------------------------------------------------------------ diff --git a/stylegan2_ae/layers/cuda/fused_bias_act.py b/stylegan2_ae/layers/cuda/fused_bias_act.py new file mode 100644 index 0000000..659f634 --- /dev/null +++ b/stylegan2_ae/layers/cuda/fused_bias_act.py @@ -0,0 +1,196 @@ +# Copyright (c) 2019, NVIDIA Corporation. All rights reserved. +# +# This work is made available under the Nvidia Source Code License-NC. +# To view a copy of this license, visit +# https://nvlabs.github.io/stylegan2/license.html + +"""Custom TensorFlow ops for efficient bias and activation.""" + +import os +import numpy as np +import tensorflow as tf +from stylegan2.layers.cuda import custom_ops +from stylegan2.utils import EasyDict + +def _get_plugin(): + return custom_ops.get_plugin(os.path.splitext(__file__)[0] + '.cu') + +#---------------------------------------------------------------------------- + +activation_funcs = { + 'linear': EasyDict(func=lambda x, **_: x, def_alpha=None, def_gain=1.0, cuda_idx=1, ref='y', zero_2nd_grad=True), + 'relu': EasyDict(func=lambda x, **_: tf.nn.relu(x), def_alpha=None, def_gain=np.sqrt(2), cuda_idx=2, ref='y', zero_2nd_grad=True), + 'lrelu': EasyDict(func=lambda x, alpha, **_: tf.nn.leaky_relu(x, alpha), def_alpha=0.2, def_gain=np.sqrt(2), cuda_idx=3, ref='y', zero_2nd_grad=True), + 'tanh': EasyDict(func=lambda x, **_: tf.nn.tanh(x), def_alpha=None, def_gain=1.0, cuda_idx=4, ref='y', zero_2nd_grad=False), + 'sigmoid': EasyDict(func=lambda x, **_: tf.nn.sigmoid(x), def_alpha=None, def_gain=1.0, cuda_idx=5, ref='y', zero_2nd_grad=False), + 'elu': EasyDict(func=lambda x, **_: tf.nn.elu(x), def_alpha=None, def_gain=1.0, cuda_idx=6, ref='y', zero_2nd_grad=False), + 'selu': EasyDict(func=lambda x, **_: tf.nn.selu(x), def_alpha=None, def_gain=1.0, cuda_idx=7, ref='y', zero_2nd_grad=False), + 'softplus': EasyDict(func=lambda x, **_: tf.nn.softplus(x), def_alpha=None, def_gain=1.0, cuda_idx=8, ref='y', zero_2nd_grad=False), + 'swish': EasyDict(func=lambda x, **_: tf.nn.sigmoid(x) * x, def_alpha=None, def_gain=np.sqrt(2), cuda_idx=9, ref='x', zero_2nd_grad=False), +} + +#---------------------------------------------------------------------------- + +def fused_bias_act(x, b=None, axis=1, act='linear', alpha=None, gain=None, impl='cuda'): + r"""Fused bias and activation function. + + Adds bias `b` to activation tensor `x`, evaluates activation function `act`, + and scales the result by `gain`. Each of the steps is optional. In most cases, + the fused op is considerably more efficient than performing the same calculation + using standard TensorFlow ops. It supports first and second order gradients, + but not third order gradients. + + Args: + x: Input activation tensor. Can have any shape, but if `b` is defined, the + dimension corresponding to `axis`, as well as the rank, must be known. + b: Bias vector, or `None` to disable. Must be a 1D tensor of the same type + as `x`. The shape must be known, and it must match the dimension of `x` + corresponding to `axis`. + axis: The dimension in `x` corresponding to the elements of `b`. + The value of `axis` is ignored if `b` is not specified. + act: Name of the activation function to evaluate, or `"linear"` to disable. + Can be e.g. `"relu"`, `"lrelu"`, `"tanh"`, `"sigmoid"`, `"swish"`, etc. + See `activation_funcs` for a full list. `None` is not allowed. + alpha: Shape parameter for the activation function, or `None` to use the default. + gain: Scaling factor for the output tensor, or `None` to use default. + See `activation_funcs` for the default scaling of each activation function. + If unsure, consider specifying `1.0`. + impl: Name of the implementation to use. Can be `"ref"` or `"cuda"` (default). + + Returns: + Tensor of the same shape and datatype as `x`. + """ + + impl_dict = { + 'ref': _fused_bias_act_ref, + 'cuda': _fused_bias_act_cuda, + } + return impl_dict[impl](x=x, b=b, axis=axis, act=act, alpha=alpha, gain=gain) + +#---------------------------------------------------------------------------- + +def _fused_bias_act_ref(x, b, axis, act, alpha, gain): + """Slow reference implementation of `fused_bias_act()` using standard TensorFlow ops.""" + + # Validate arguments. + x = tf.convert_to_tensor(x) + b = tf.convert_to_tensor(b) if b is not None else tf.constant([], dtype=x.dtype) + act_spec = activation_funcs[act] + assert b.shape.rank == 1 and (b.shape[0] == 0 or b.shape[0] == x.shape[axis]) + assert b.shape[0] == 0 or 0 <= axis < x.shape.rank + if alpha is None: + alpha = act_spec.def_alpha + if gain is None: + gain = act_spec.def_gain + + # Add bias. + if b.shape[0] != 0: + x += tf.reshape(b, [-1 if i == axis else 1 for i in range(x.shape.rank)]) + + # Evaluate activation function. + x = act_spec.func(x, alpha=alpha) + + # Scale by gain. + if gain != 1: + x *= gain + return x + +#---------------------------------------------------------------------------- + +def _fused_bias_act_cuda(x, b, axis, act, alpha, gain): + """Fast CUDA implementation of `fused_bias_act()` using custom ops.""" + + # Validate arguments. + x = tf.convert_to_tensor(x) + empty_tensor = tf.constant([], dtype=x.dtype) + b = tf.convert_to_tensor(b) if b is not None else empty_tensor + act_spec = activation_funcs[act] + assert b.shape.rank == 1 and (b.shape[0] == 0 or b.shape[0] == x.shape[axis]) + assert b.shape[0] == 0 or 0 <= axis < x.shape.rank + if alpha is None: + alpha = act_spec.def_alpha + if gain is None: + gain = act_spec.def_gain + + # Special cases. + if act == 'linear' and b is None and gain == 1.0: + return x + if act_spec.cuda_idx is None: + return _fused_bias_act_ref(x=x, b=b, axis=axis, act=act, alpha=alpha, gain=gain) + + # CUDA kernel. + cuda_kernel = _get_plugin().fused_bias_act + cuda_kwargs = dict(axis=axis, act=act_spec.cuda_idx, alpha=alpha, gain=gain) + + # Forward pass: y = func(x, b). + def func_y(x, b): + y = cuda_kernel(x=x, b=b, ref=empty_tensor, grad=0, **cuda_kwargs) + y.set_shape(x.shape) + return y + + # Backward pass: dx, db = grad(dy, x, y) + def grad_dx(dy, x, y): + ref = {'x': x, 'y': y}[act_spec.ref] + dx = cuda_kernel(x=dy, b=empty_tensor, ref=ref, grad=1, **cuda_kwargs) + dx.set_shape(x.shape) + return dx + def grad_db(dx): + if b.shape[0] == 0: + return empty_tensor + db = dx + if axis < x.shape.rank - 1: + db = tf.reduce_sum(db, list(range(axis + 1, x.shape.rank))) + if axis > 0: + db = tf.reduce_sum(db, list(range(axis))) + db.set_shape(b.shape) + return db + + # Second order gradients: d_dy, d_x = grad2(d_dx, d_db, x, y) + def grad2_d_dy(d_dx, d_db, x, y): + ref = {'x': x, 'y': y}[act_spec.ref] + d_dy = cuda_kernel(x=d_dx, b=d_db, ref=ref, grad=1, **cuda_kwargs) + d_dy.set_shape(x.shape) + return d_dy + def grad2_d_x(d_dx, d_db, x, y): + ref = {'x': x, 'y': y}[act_spec.ref] + d_x = cuda_kernel(x=d_dx, b=d_db, ref=ref, grad=2, **cuda_kwargs) + d_x.set_shape(x.shape) + return d_x + + # Fast version for piecewise-linear activation funcs. + @tf.custom_gradient + def func_zero_2nd_grad(x, b): + y = func_y(x, b) + @tf.custom_gradient + def grad(dy): + dx = grad_dx(dy, x, y) + db = grad_db(dx) + def grad2(d_dx, d_db): + d_dy = grad2_d_dy(d_dx, d_db, x, y) + return d_dy + return (dx, db), grad2 + return y, grad + + # Slow version for general activation funcs. + @tf.custom_gradient + def func_nonzero_2nd_grad(x, b): + y = func_y(x, b) + def grad_wrap(dy): + @tf.custom_gradient + def grad_impl(dy, x): + dx = grad_dx(dy, x, y) + db = grad_db(dx) + def grad2(d_dx, d_db): + d_dy = grad2_d_dy(d_dx, d_db, x, y) + d_x = grad2_d_x(d_dx, d_db, x, y) + return d_dy, d_x + return (dx, db), grad2 + return grad_impl(dy, x) + return y, grad_wrap + + # Which version to use? + if act_spec.zero_2nd_grad: + return func_zero_2nd_grad(x, b) + return func_nonzero_2nd_grad(x, b) + +#---------------------------------------------------------------------------- diff --git a/stylegan2_ae/layers/cuda/upfirdn_2d.cu b/stylegan2_ae/layers/cuda/upfirdn_2d.cu new file mode 100644 index 0000000..b97ef36 --- /dev/null +++ b/stylegan2_ae/layers/cuda/upfirdn_2d.cu @@ -0,0 +1,326 @@ +// Copyright (c) 2019, NVIDIA Corporation. All rights reserved. +// +// This work is made available under the Nvidia Source Code License-NC. +// To view a copy of this license, visit +// https://nvlabs.github.io/stylegan2/license.html + +#define EIGEN_USE_GPU +#define __CUDA_INCLUDE_COMPILER_INTERNAL_HEADERS__ +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/shape_inference.h" +#include + +using namespace tensorflow; +using namespace tensorflow::shape_inference; + +//------------------------------------------------------------------------ +// Helpers. + +#define OP_CHECK_CUDA_ERROR(CTX, CUDA_CALL) do { cudaError_t err = CUDA_CALL; OP_REQUIRES(CTX, err == cudaSuccess, errors::Internal(cudaGetErrorName(err))); } while (false) + +static __host__ __device__ __forceinline__ int floorDiv(int a, int b) +{ + int c = a / b; + if (c * b > a) + c--; + return c; +} + +//------------------------------------------------------------------------ +// CUDA kernel params. + +template +struct UpFirDn2DKernelParams +{ + const T* x; // [majorDim, inH, inW, minorDim] + const T* k; // [kernelH, kernelW] + T* y; // [majorDim, outH, outW, minorDim] + + int upx; + int upy; + int downx; + int downy; + int padx0; + int padx1; + int pady0; + int pady1; + + int majorDim; + int inH; + int inW; + int minorDim; + int kernelH; + int kernelW; + int outH; + int outW; + int loopMajor; + int loopX; +}; + +//------------------------------------------------------------------------ +// General CUDA implementation for large filter kernels. + +template +static __global__ void UpFirDn2DKernel_large(const UpFirDn2DKernelParams p) +{ + // Calculate thread index. + int minorIdx = blockIdx.x * blockDim.x + threadIdx.x; + int outY = minorIdx / p.minorDim; + minorIdx -= outY * p.minorDim; + int outXBase = blockIdx.y * p.loopX * blockDim.y + threadIdx.y; + int majorIdxBase = blockIdx.z * p.loopMajor; + if (outXBase >= p.outW || outY >= p.outH || majorIdxBase >= p.majorDim) + return; + + // Setup Y receptive field. + int midY = outY * p.downy + p.upy - 1 - p.pady0; + int inY = min(max(floorDiv(midY, p.upy), 0), p.inH); + int h = min(max(floorDiv(midY + p.kernelH, p.upy), 0), p.inH) - inY; + int kernelY = midY + p.kernelH - (inY + 1) * p.upy; + + // Loop over majorDim and outX. + for (int loopMajor = 0, majorIdx = majorIdxBase; loopMajor < p.loopMajor && majorIdx < p.majorDim; loopMajor++, majorIdx++) + for (int loopX = 0, outX = outXBase; loopX < p.loopX && outX < p.outW; loopX++, outX += blockDim.y) + { + // Setup X receptive field. + int midX = outX * p.downx + p.upx - 1 - p.padx0; + int inX = min(max(floorDiv(midX, p.upx), 0), p.inW); + int w = min(max(floorDiv(midX + p.kernelW, p.upx), 0), p.inW) - inX; + int kernelX = midX + p.kernelW - (inX + 1) * p.upx; + + // Initialize pointers. + const T* xp = &p.x[((majorIdx * p.inH + inY) * p.inW + inX) * p.minorDim + minorIdx]; + const T* kp = &p.k[kernelY * p.kernelW + kernelX]; + int xpx = p.minorDim; + int kpx = -p.upx; + int xpy = p.inW * p.minorDim; + int kpy = -p.upy * p.kernelW; + + // Inner loop. + float v = 0.0f; + for (int y = 0; y < h; y++) + { + for (int x = 0; x < w; x++) + { + v += (float)(*xp) * (float)(*kp); + xp += xpx; + kp += kpx; + } + xp += xpy - w * xpx; + kp += kpy - w * kpx; + } + + // Store result. + p.y[((majorIdx * p.outH + outY) * p.outW + outX) * p.minorDim + minorIdx] = (T)v; + } +} + +//------------------------------------------------------------------------ +// Specialized CUDA implementation for small filter kernels. + +template +static __global__ void UpFirDn2DKernel_small(const UpFirDn2DKernelParams p) +{ + //assert(kernelW % upx == 0); + //assert(kernelH % upy == 0); + const int tileInW = ((tileOutW - 1) * downx + kernelW - 1) / upx + 1; + const int tileInH = ((tileOutH - 1) * downy + kernelH - 1) / upy + 1; + __shared__ volatile float sk[kernelH][kernelW]; + __shared__ volatile float sx[tileInH][tileInW]; + + // Calculate tile index. + int minorIdx = blockIdx.x; + int tileOutY = minorIdx / p.minorDim; + minorIdx -= tileOutY * p.minorDim; + tileOutY *= tileOutH; + int tileOutXBase = blockIdx.y * p.loopX * tileOutW; + int majorIdxBase = blockIdx.z * p.loopMajor; + if (tileOutXBase >= p.outW | tileOutY >= p.outH | majorIdxBase >= p.majorDim) + return; + + // Load filter kernel (flipped). + for (int tapIdx = threadIdx.x; tapIdx < kernelH * kernelW; tapIdx += blockDim.x) + { + int ky = tapIdx / kernelW; + int kx = tapIdx - ky * kernelW; + float v = 0.0f; + if (kx < p.kernelW & ky < p.kernelH) + v = (float)p.k[(p.kernelH - 1 - ky) * p.kernelW + (p.kernelW - 1 - kx)]; + sk[ky][kx] = v; + } + + // Loop over majorDim and outX. + for (int loopMajor = 0, majorIdx = majorIdxBase; loopMajor < p.loopMajor & majorIdx < p.majorDim; loopMajor++, majorIdx++) + for (int loopX = 0, tileOutX = tileOutXBase; loopX < p.loopX & tileOutX < p.outW; loopX++, tileOutX += tileOutW) + { + // Load input pixels. + int tileMidX = tileOutX * downx + upx - 1 - p.padx0; + int tileMidY = tileOutY * downy + upy - 1 - p.pady0; + int tileInX = floorDiv(tileMidX, upx); + int tileInY = floorDiv(tileMidY, upy); + __syncthreads(); + for (int inIdx = threadIdx.x; inIdx < tileInH * tileInW; inIdx += blockDim.x) + { + int relInY = inIdx / tileInW; + int relInX = inIdx - relInY * tileInW; + int inX = relInX + tileInX; + int inY = relInY + tileInY; + float v = 0.0f; + if (inX >= 0 & inY >= 0 & inX < p.inW & inY < p.inH) + v = (float)p.x[((majorIdx * p.inH + inY) * p.inW + inX) * p.minorDim + minorIdx]; + sx[relInY][relInX] = v; + } + + // Loop over output pixels. + __syncthreads(); + for (int outIdx = threadIdx.x; outIdx < tileOutH * tileOutW; outIdx += blockDim.x) + { + int relOutY = outIdx / tileOutW; + int relOutX = outIdx - relOutY * tileOutW; + int outX = relOutX + tileOutX; + int outY = relOutY + tileOutY; + + // Setup receptive field. + int midX = tileMidX + relOutX * downx; + int midY = tileMidY + relOutY * downy; + int inX = floorDiv(midX, upx); + int inY = floorDiv(midY, upy); + int relInX = inX - tileInX; + int relInY = inY - tileInY; + int kernelX = (inX + 1) * upx - midX - 1; // flipped + int kernelY = (inY + 1) * upy - midY - 1; // flipped + + // Inner loop. + float v = 0.0f; + #pragma unroll + for (int y = 0; y < kernelH / upy; y++) + #pragma unroll + for (int x = 0; x < kernelW / upx; x++) + v += sx[relInY + y][relInX + x] * sk[kernelY + y * upy][kernelX + x * upx]; + + // Store result. + if (outX < p.outW & outY < p.outH) + p.y[((majorIdx * p.outH + outY) * p.outW + outX) * p.minorDim + minorIdx] = (T)v; + } + } +} + +//------------------------------------------------------------------------ +// TensorFlow op. + +template +struct UpFirDn2DOp : public OpKernel +{ + UpFirDn2DKernelParams m_attribs; + + UpFirDn2DOp(OpKernelConstruction* ctx) : OpKernel(ctx) + { + memset(&m_attribs, 0, sizeof(m_attribs)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("upx", &m_attribs.upx)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("upy", &m_attribs.upy)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("downx", &m_attribs.downx)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("downy", &m_attribs.downy)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("padx0", &m_attribs.padx0)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("padx1", &m_attribs.padx1)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("pady0", &m_attribs.pady0)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("pady1", &m_attribs.pady1)); + OP_REQUIRES(ctx, m_attribs.upx >= 1 && m_attribs.upy >= 1, errors::InvalidArgument("upx and upy must be at least 1x1")); + OP_REQUIRES(ctx, m_attribs.downx >= 1 && m_attribs.downy >= 1, errors::InvalidArgument("downx and downy must be at least 1x1")); + } + + void Compute(OpKernelContext* ctx) + { + UpFirDn2DKernelParams p = m_attribs; + cudaStream_t stream = ctx->eigen_device().stream(); + + const Tensor& x = ctx->input(0); // [majorDim, inH, inW, minorDim] + const Tensor& k = ctx->input(1); // [kernelH, kernelW] + p.x = x.flat().data(); + p.k = k.flat().data(); + OP_REQUIRES(ctx, x.dims() == 4, errors::InvalidArgument("input must have rank 4")); + OP_REQUIRES(ctx, k.dims() == 2, errors::InvalidArgument("kernel must have rank 2")); + OP_REQUIRES(ctx, x.NumElements() <= kint32max, errors::InvalidArgument("input too large")); + OP_REQUIRES(ctx, k.NumElements() <= kint32max, errors::InvalidArgument("kernel too large")); + + p.majorDim = (int)x.dim_size(0); + p.inH = (int)x.dim_size(1); + p.inW = (int)x.dim_size(2); + p.minorDim = (int)x.dim_size(3); + p.kernelH = (int)k.dim_size(0); + p.kernelW = (int)k.dim_size(1); + OP_REQUIRES(ctx, p.kernelW >= 1 && p.kernelH >= 1, errors::InvalidArgument("kernel must be at least 1x1")); + + p.outW = (p.inW * p.upx + p.padx0 + p.padx1 - p.kernelW + p.downx) / p.downx; + p.outH = (p.inH * p.upy + p.pady0 + p.pady1 - p.kernelH + p.downy) / p.downy; + OP_REQUIRES(ctx, p.outW >= 1 && p.outH >= 1, errors::InvalidArgument("output must be at least 1x1")); + + Tensor* y = NULL; // [majorDim, outH, outW, minorDim] + TensorShape ys; + ys.AddDim(p.majorDim); + ys.AddDim(p.outH); + ys.AddDim(p.outW); + ys.AddDim(p.minorDim); + OP_REQUIRES_OK(ctx, ctx->allocate_output(0, ys, &y)); + p.y = y->flat().data(); + OP_REQUIRES(ctx, y->NumElements() <= kint32max, errors::InvalidArgument("output too large")); + + // Choose CUDA kernel to use. + void* cudaKernel = (void*)UpFirDn2DKernel_large; + int tileOutW = -1; + int tileOutH = -1; + if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 7 && p.kernelH <= 7) { cudaKernel = (void*)UpFirDn2DKernel_small; tileOutW = 64; tileOutH = 16; } + if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 6 && p.kernelH <= 6) { cudaKernel = (void*)UpFirDn2DKernel_small; tileOutW = 64; tileOutH = 16; } + if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 5 && p.kernelH <= 5) { cudaKernel = (void*)UpFirDn2DKernel_small; tileOutW = 64; tileOutH = 16; } + if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 4 && p.kernelH <= 4) { cudaKernel = (void*)UpFirDn2DKernel_small; tileOutW = 64; tileOutH = 16; } + if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 3 && p.kernelH <= 3) { cudaKernel = (void*)UpFirDn2DKernel_small; tileOutW = 64; tileOutH = 16; } + if (p.upx == 2 && p.upy == 2 && p.downx == 1 && p.downy == 1 && p.kernelW <= 8 && p.kernelH <= 8) { cudaKernel = (void*)UpFirDn2DKernel_small; tileOutW = 64; tileOutH = 16; } + if (p.upx == 2 && p.upy == 2 && p.downx == 1 && p.downy == 1 && p.kernelW <= 6 && p.kernelH <= 6) { cudaKernel = (void*)UpFirDn2DKernel_small; tileOutW = 64; tileOutH = 16; } + if (p.upx == 2 && p.upy == 2 && p.downx == 1 && p.downy == 1 && p.kernelW <= 4 && p.kernelH <= 4) { cudaKernel = (void*)UpFirDn2DKernel_small; tileOutW = 64; tileOutH = 16; } + if (p.upx == 2 && p.upy == 2 && p.downx == 1 && p.downy == 1 && p.kernelW <= 2 && p.kernelH <= 2) { cudaKernel = (void*)UpFirDn2DKernel_small; tileOutW = 64; tileOutH = 16; } + if (p.upx == 1 && p.upy == 1 && p.downx == 2 && p.downy == 2 && p.kernelW <= 8 && p.kernelH <= 8) { cudaKernel = (void*)UpFirDn2DKernel_small; tileOutW = 32; tileOutH = 8; } + if (p.upx == 1 && p.upy == 1 && p.downx == 2 && p.downy == 2 && p.kernelW <= 6 && p.kernelH <= 6) { cudaKernel = (void*)UpFirDn2DKernel_small; tileOutW = 32; tileOutH = 8; } + if (p.upx == 1 && p.upy == 1 && p.downx == 2 && p.downy == 2 && p.kernelW <= 4 && p.kernelH <= 4) { cudaKernel = (void*)UpFirDn2DKernel_small; tileOutW = 32; tileOutH = 8; } + if (p.upx == 1 && p.upy == 1 && p.downx == 2 && p.downy == 2 && p.kernelW <= 2 && p.kernelH <= 2) { cudaKernel = (void*)UpFirDn2DKernel_small; tileOutW = 32; tileOutH = 8; } + + // Choose launch params. + dim3 blockSize; + dim3 gridSize; + if (tileOutW > 0 && tileOutH > 0) // small + { + p.loopMajor = (p.majorDim - 1) / 16384 + 1; + p.loopX = 1; + blockSize = dim3(32 * 8, 1, 1); + gridSize = dim3(((p.outH - 1) / tileOutH + 1) * p.minorDim, (p.outW - 1) / (p.loopX * tileOutW) + 1, (p.majorDim - 1) / p.loopMajor + 1); + } + else // large + { + p.loopMajor = (p.majorDim - 1) / 16384 + 1; + p.loopX = 4; + blockSize = dim3(4, 32, 1); + gridSize = dim3((p.outH * p.minorDim - 1) / blockSize.x + 1, (p.outW - 1) / (p.loopX * blockSize.y) + 1, (p.majorDim - 1) / p.loopMajor + 1); + } + + // Launch CUDA kernel. + void* args[] = {&p}; + OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel(cudaKernel, gridSize, blockSize, args, 0, stream)); + } +}; + +REGISTER_OP("UpFirDn2D") + .Input ("x: T") + .Input ("k: T") + .Output ("y: T") + .Attr ("T: {float, half}") + .Attr ("upx: int = 1") + .Attr ("upy: int = 1") + .Attr ("downx: int = 1") + .Attr ("downy: int = 1") + .Attr ("padx0: int = 0") + .Attr ("padx1: int = 0") + .Attr ("pady0: int = 0") + .Attr ("pady1: int = 0"); +REGISTER_KERNEL_BUILDER(Name("UpFirDn2D").Device(DEVICE_GPU).TypeConstraint("T"), UpFirDn2DOp); +REGISTER_KERNEL_BUILDER(Name("UpFirDn2D").Device(DEVICE_GPU).TypeConstraint("T"), UpFirDn2DOp); + +//------------------------------------------------------------------------ diff --git a/stylegan2_ae/layers/cuda/upfirdn_2d.py b/stylegan2_ae/layers/cuda/upfirdn_2d.py new file mode 100644 index 0000000..4194915 --- /dev/null +++ b/stylegan2_ae/layers/cuda/upfirdn_2d.py @@ -0,0 +1,364 @@ +# Copyright (c) 2019, NVIDIA Corporation. All rights reserved. +# +# This work is made available under the Nvidia Source Code License-NC. +# To view a copy of this license, visit +# https://nvlabs.github.io/stylegan2/license.html + +"""Custom TensorFlow ops for efficient resampling of 2D images.""" + +import os +import numpy as np +import tensorflow as tf +from stylegan2.layers.cuda import custom_ops + +def _get_plugin(): + return custom_ops.get_plugin(os.path.splitext(__file__)[0] + '.cu') + +#---------------------------------------------------------------------------- + +def upfirdn_2d(x, k, upx=1, upy=1, downx=1, downy=1, padx0=0, padx1=0, pady0=0, pady1=0, impl='cuda'): + r"""Pad, upsample, FIR filter, and downsample a batch of 2D images. + + Accepts a batch of 2D images of the shape `[majorDim, inH, inW, minorDim]` + and performs the following operations for each image, batched across + `majorDim` and `minorDim`: + + 1. Pad the image with zeros by the specified number of pixels on each side + (`padx0`, `padx1`, `pady0`, `pady1`). Specifying a negative value + corresponds to cropping the image. + + 2. Upsample the image by inserting the zeros after each pixel (`upx`, `upy`). + + 3. Convolve the image with the specified 2D FIR filter (`k`), shrinking the + image so that the footprint of all output pixels lies within the input image. + + 4. Downsample the image by throwing away pixels (`downx`, `downy`). + + This sequence of operations bears close resemblance to scipy.signal.upfirdn(). + The fused op is considerably more efficient than performing the same calculation + using standard TensorFlow ops. It supports gradients of arbitrary order. + + Args: + x: Input tensor of the shape `[majorDim, inH, inW, minorDim]`. + k: 2D FIR filter of the shape `[firH, firW]`. + upx: Integer upsampling factor along the X-axis (default: 1). + upy: Integer upsampling factor along the Y-axis (default: 1). + downx: Integer downsampling factor along the X-axis (default: 1). + downy: Integer downsampling factor along the Y-axis (default: 1). + padx0: Number of pixels to pad on the left side (default: 0). + padx1: Number of pixels to pad on the right side (default: 0). + pady0: Number of pixels to pad on the top side (default: 0). + pady1: Number of pixels to pad on the bottom side (default: 0). + impl: Name of the implementation to use. Can be `"ref"` or `"cuda"` (default). + + Returns: + Tensor of the shape `[majorDim, outH, outW, minorDim]`, and same datatype as `x`. + """ + + impl_dict = { + 'ref': _upfirdn_2d_ref, + 'cuda': _upfirdn_2d_cuda, + } + return impl_dict[impl](x=x, k=k, upx=upx, upy=upy, downx=downx, downy=downy, padx0=padx0, padx1=padx1, pady0=pady0, pady1=pady1) + +#---------------------------------------------------------------------------- + +def _upfirdn_2d_ref(x, k, upx, upy, downx, downy, padx0, padx1, pady0, pady1): + """Slow reference implementation of `upfirdn_2d()` using standard TensorFlow ops.""" + + x = tf.convert_to_tensor(x) + k = np.asarray(k, dtype=np.float32) + assert x.shape.rank == 4 + inH = x.shape[1] + inW = x.shape[2] + minorDim = _shape(x, 3) + kernelH, kernelW = k.shape + assert inW >= 1 and inH >= 1 + assert kernelW >= 1 and kernelH >= 1 + assert isinstance(upx, int) and isinstance(upy, int) + assert isinstance(downx, int) and isinstance(downy, int) + assert isinstance(padx0, int) and isinstance(padx1, int) + assert isinstance(pady0, int) and isinstance(pady1, int) + + # Upsample (insert zeros). + x = tf.reshape(x, [-1, inH, 1, inW, 1, minorDim]) + x = tf.pad(x, [[0, 0], [0, 0], [0, upy - 1], [0, 0], [0, upx - 1], [0, 0]]) + x = tf.reshape(x, [-1, inH * upy, inW * upx, minorDim]) + + # Pad (crop if negative). + x = tf.pad(x, [[0, 0], [max(pady0, 0), max(pady1, 0)], [max(padx0, 0), max(padx1, 0)], [0, 0]]) + x = x[:, max(-pady0, 0) : x.shape[1] - max(-pady1, 0), max(-padx0, 0) : x.shape[2] - max(-padx1, 0), :] + + # Convolve with filter. + x = tf.transpose(x, [0, 3, 1, 2]) + x = tf.reshape(x, [-1, 1, inH * upy + pady0 + pady1, inW * upx + padx0 + padx1]) + w = tf.constant(k[::-1, ::-1, np.newaxis, np.newaxis], dtype=x.dtype) + x = tf.nn.conv2d(x, w, strides=[1,1,1,1], padding='VALID', data_format='NCHW') + x = tf.reshape(x, [-1, minorDim, inH * upy + pady0 + pady1 - kernelH + 1, inW * upx + padx0 + padx1 - kernelW + 1]) + x = tf.transpose(x, [0, 2, 3, 1]) + + # Downsample (throw away pixels). + return x[:, ::downy, ::downx, :] + +#---------------------------------------------------------------------------- + +def _upfirdn_2d_cuda(x, k, upx, upy, downx, downy, padx0, padx1, pady0, pady1): + """Fast CUDA implementation of `upfirdn_2d()` using custom ops.""" + + x = tf.convert_to_tensor(x) + k = np.asarray(k, dtype=np.float32) + majorDim, inH, inW, minorDim = x.shape.as_list() + kernelH, kernelW = k.shape + assert inW >= 1 and inH >= 1 + assert kernelW >= 1 and kernelH >= 1 + assert isinstance(upx, int) and isinstance(upy, int) + assert isinstance(downx, int) and isinstance(downy, int) + assert isinstance(padx0, int) and isinstance(padx1, int) + assert isinstance(pady0, int) and isinstance(pady1, int) + + outW = (inW * upx + padx0 + padx1 - kernelW) // downx + 1 + outH = (inH * upy + pady0 + pady1 - kernelH) // downy + 1 + assert outW >= 1 and outH >= 1 + + kc = tf.constant(k, dtype=x.dtype) + gkc = tf.constant(k[::-1, ::-1], dtype=x.dtype) + gpadx0 = kernelW - padx0 - 1 + gpady0 = kernelH - pady0 - 1 + gpadx1 = inW * upx - outW * downx + padx0 - upx + 1 + gpady1 = inH * upy - outH * downy + pady0 - upy + 1 + + @tf.custom_gradient + def func(x): + y = _get_plugin().up_fir_dn2d(x=x, k=kc, upx=upx, upy=upy, downx=downx, downy=downy, padx0=padx0, padx1=padx1, pady0=pady0, pady1=pady1) + y.set_shape([majorDim, outH, outW, minorDim]) + @tf.custom_gradient + def grad(dy): + dx = _get_plugin().up_fir_dn2d(x=dy, k=gkc, upx=downx, upy=downy, downx=upx, downy=upy, padx0=gpadx0, padx1=gpadx1, pady0=gpady0, pady1=gpady1) + dx.set_shape([majorDim, inH, inW, minorDim]) + return dx, func + return y, grad + return func(x) + +#---------------------------------------------------------------------------- + +def filter_2d(x, k, gain=1, data_format='NCHW', impl='cuda'): + r"""Filter a batch of 2D images with the given FIR filter. + + Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]` + and filters each image with the given filter. The filter is normalized so that + if the input pixels are constant, they will be scaled by the specified `gain`. + Pixels outside the image are assumed to be zero. + + Args: + x: Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`. + k: FIR filter of the shape `[firH, firW]` or `[firN]` (separable). + gain: Scaling factor for signal magnitude (default: 1.0). + data_format: `'NCHW'` or `'NHWC'` (default: `'NCHW'`). + impl: Name of the implementation to use. Can be `"ref"` or `"cuda"` (default). + + Returns: + Tensor of the same shape and datatype as `x`. + """ + + k = _setup_kernel(k) * gain + p = k.shape[0] - 1 + return _simple_upfirdn_2d(x, k, pad0=(p+1)//2, pad1=p//2, data_format=data_format, impl=impl) + +#---------------------------------------------------------------------------- + +def upsample_2d(x, k=None, factor=2, gain=1, data_format='NCHW', impl='cuda'): + r"""Upsample a batch of 2D images with the given filter. + + Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]` + and upsamples each image with the given filter. The filter is normalized so that + if the input pixels are constant, they will be scaled by the specified `gain`. + Pixels outside the image are assumed to be zero, and the filter is padded with + zeros so that its shape is a multiple of the upsampling factor. + + Args: + x: Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`. + k: FIR filter of the shape `[firH, firW]` or `[firN]` (separable). + The default is `[1] * factor`, which corresponds to nearest-neighbor + upsampling. + factor: Integer upsampling factor (default: 2). + gain: Scaling factor for signal magnitude (default: 1.0). + data_format: `'NCHW'` or `'NHWC'` (default: `'NCHW'`). + impl: Name of the implementation to use. Can be `"ref"` or `"cuda"` (default). + + Returns: + Tensor of the shape `[N, C, H * factor, W * factor]` or + `[N, H * factor, W * factor, C]`, and same datatype as `x`. + """ + + assert isinstance(factor, int) and factor >= 1 + if k is None: + k = [1] * factor + k = _setup_kernel(k) * (gain * (factor ** 2)) + p = k.shape[0] - factor + return _simple_upfirdn_2d(x, k, up=factor, pad0=(p+1)//2+factor-1, pad1=p//2, data_format=data_format, impl=impl) + +#---------------------------------------------------------------------------- + +def downsample_2d(x, k=None, factor=2, gain=1, data_format='NCHW', impl='cuda'): + r"""Downsample a batch of 2D images with the given filter. + + Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]` + and downsamples each image with the given filter. The filter is normalized so that + if the input pixels are constant, they will be scaled by the specified `gain`. + Pixels outside the image are assumed to be zero, and the filter is padded with + zeros so that its shape is a multiple of the downsampling factor. + + Args: + x: Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`. + k: FIR filter of the shape `[firH, firW]` or `[firN]` (separable). + The default is `[1] * factor`, which corresponds to average pooling. + factor: Integer downsampling factor (default: 2). + gain: Scaling factor for signal magnitude (default: 1.0). + data_format: `'NCHW'` or `'NHWC'` (default: `'NCHW'`). + impl: Name of the implementation to use. Can be `"ref"` or `"cuda"` (default). + + Returns: + Tensor of the shape `[N, C, H // factor, W // factor]` or + `[N, H // factor, W // factor, C]`, and same datatype as `x`. + """ + + assert isinstance(factor, int) and factor >= 1 + if k is None: + k = [1] * factor + k = _setup_kernel(k) * gain + p = k.shape[0] - factor + return _simple_upfirdn_2d(x, k, down=factor, pad0=(p+1)//2, pad1=p//2, data_format=data_format, impl=impl) + +#---------------------------------------------------------------------------- + +def upsample_conv_2d(x, w, k=None, factor=2, gain=1, data_format='NCHW', impl='cuda'): + r"""Fused `upsample_2d()` followed by `tf.nn.conv2d()`. + + Padding is performed only once at the beginning, not between the operations. + The fused op is considerably more efficient than performing the same calculation + using standard TensorFlow ops. It supports gradients of arbitrary order. + + Args: + x: Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`. + w: Weight tensor of the shape `[filterH, filterW, inChannels, outChannels]`. + Grouped convolution can be performed by `inChannels = x.shape[0] // numGroups`. + k: FIR filter of the shape `[firH, firW]` or `[firN]` (separable). + The default is `[1] * factor`, which corresponds to nearest-neighbor + upsampling. + factor: Integer upsampling factor (default: 2). + gain: Scaling factor for signal magnitude (default: 1.0). + data_format: `'NCHW'` or `'NHWC'` (default: `'NCHW'`). + impl: Name of the implementation to use. Can be `"ref"` or `"cuda"` (default). + + Returns: + Tensor of the shape `[N, C, H * factor, W * factor]` or + `[N, H * factor, W * factor, C]`, and same datatype as `x`. + """ + + assert isinstance(factor, int) and factor >= 1 + + # Check weight shape. + w = tf.convert_to_tensor(w) + assert w.shape.rank == 4 + convH = w.shape[0] + convW = w.shape[1] + inC = _shape(w, 2) + outC = _shape(w, 3) + assert convW == convH + + # Setup filter kernel. + if k is None: + k = [1] * factor + k = _setup_kernel(k) * (gain * (factor ** 2)) + p = (k.shape[0] - factor) - (convW - 1) + + # Determine data dimensions. + if data_format == 'NCHW': + stride = [1, 1, factor, factor] + output_shape = [_shape(x, 0), outC, (_shape(x, 2) - 1) * factor + convH, (_shape(x, 3) - 1) * factor + convW] + num_groups = _shape(x, 1) // inC + else: + stride = [1, factor, factor, 1] + output_shape = [_shape(x, 0), (_shape(x, 1) - 1) * factor + convH, (_shape(x, 2) - 1) * factor + convW, outC] + num_groups = _shape(x, 3) // inC + + # Transpose weights. + w = tf.reshape(w, [convH, convW, inC, num_groups, -1]) + w = tf.transpose(w[::-1, ::-1], [0, 1, 4, 3, 2]) + w = tf.reshape(w, [convH, convW, -1, num_groups * inC]) + + # Execute. + x = tf.nn.conv2d_transpose(x, w, output_shape=output_shape, strides=stride, padding='VALID', data_format=data_format) + return _simple_upfirdn_2d(x, k, pad0=(p+1)//2+factor-1, pad1=p//2+1, data_format=data_format, impl=impl) + +#---------------------------------------------------------------------------- + +def conv_downsample_2d(x, w, k=None, factor=2, gain=1, data_format='NCHW', impl='cuda'): + r"""Fused `tf.nn.conv2d()` followed by `downsample_2d()`. + + Padding is performed only once at the beginning, not between the operations. + The fused op is considerably more efficient than performing the same calculation + using standard TensorFlow ops. It supports gradients of arbitrary order. + + Args: + x: Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`. + w: Weight tensor of the shape `[filterH, filterW, inChannels, outChannels]`. + Grouped convolution can be performed by `inChannels = x.shape[0] // numGroups`. + k: FIR filter of the shape `[firH, firW]` or `[firN]` (separable). + The default is `[1] * factor`, which corresponds to average pooling. + factor: Integer downsampling factor (default: 2). + gain: Scaling factor for signal magnitude (default: 1.0). + data_format: `'NCHW'` or `'NHWC'` (default: `'NCHW'`). + impl: Name of the implementation to use. Can be `"ref"` or `"cuda"` (default). + + Returns: + Tensor of the shape `[N, C, H // factor, W // factor]` or + `[N, H // factor, W // factor, C]`, and same datatype as `x`. + """ + + assert isinstance(factor, int) and factor >= 1 + w = tf.convert_to_tensor(w) + convH, convW, _inC, _outC = w.shape.as_list() + assert convW == convH + if k is None: + k = [1] * factor + k = _setup_kernel(k) * gain + p = (k.shape[0] - factor) + (convW - 1) + if data_format == 'NCHW': + s = [1, 1, factor, factor] + else: + s = [1, factor, factor, 1] + x = _simple_upfirdn_2d(x, k, pad0=(p+1)//2, pad1=p//2, data_format=data_format, impl=impl) + return tf.nn.conv2d(x, w, strides=s, padding='VALID', data_format=data_format) + +#---------------------------------------------------------------------------- +# Internal helper funcs. + +def _shape(tf_expr, dim_idx): + if tf_expr.shape.rank is not None: + dim = tf_expr.shape[dim_idx] + if dim is not None: + return dim + return tf.shape(tf_expr)[dim_idx] + +def _setup_kernel(k): + k = np.asarray(k, dtype=np.float32) + if k.ndim == 1: + k = np.outer(k, k) + k /= np.sum(k) + assert k.ndim == 2 + assert k.shape[0] == k.shape[1] + return k + +def _simple_upfirdn_2d(x, k, up=1, down=1, pad0=0, pad1=0, data_format='NCHW', impl='cuda'): + assert data_format in ['NCHW', 'NHWC'] + assert x.shape.rank == 4 + y = x + if data_format == 'NCHW': + y = tf.reshape(y, [-1, _shape(y, 2), _shape(y, 3), 1]) + y = upfirdn_2d(y, k, upx=up, upy=up, downx=down, downy=down, padx0=pad0, padx1=pad1, pady0=pad0, pady1=pad1, impl=impl) + if data_format == 'NCHW': + y = tf.reshape(y, [-1, _shape(x, 1), _shape(y, 1), _shape(y, 2)]) + return y + +#---------------------------------------------------------------------------- diff --git a/stylegan2_ae/layers/cuda/upfirdn_2d_v2.py b/stylegan2_ae/layers/cuda/upfirdn_2d_v2.py new file mode 100644 index 0000000..69eec15 --- /dev/null +++ b/stylegan2_ae/layers/cuda/upfirdn_2d_v2.py @@ -0,0 +1,146 @@ +import os +import numpy as np +import tensorflow as tf +from stylegan2.layers.cuda import custom_ops + + +def _get_plugin(): + loc = os.path.dirname(os.path.abspath(__file__)) + cu_fn = 'upfirdn_2d.cu' + return custom_ops.get_plugin(os.path.join(loc, cu_fn)) + + +def _setup_kernel(k): + k = np.asarray(k, dtype=np.float32) + if k.ndim == 1: + k = np.outer(k, k) + k /= np.sum(k) + assert k.ndim == 2 + assert k.shape[0] == k.shape[1] + return k + + +def compute_paddings(resample_kernel, convW, up, down, is_conv, factor=2, gain=1): + assert not (up and down) + + k = [1] * factor if resample_kernel is None else resample_kernel + if up: + k = _setup_kernel(k) * (gain * (factor ** 2)) + if is_conv: + p = (k.shape[0] - factor) - (convW - 1) + pad0 = (p + 1) // 2 + factor - 1 + pad1 = p // 2 + 1 + else: + p = k.shape[0] - factor + pad0 = (p + 1) // 2 + factor - 1 + pad1 = p // 2 + elif down: + k = _setup_kernel(k) * gain + if is_conv: + p = (k.shape[0] - factor) + (convW - 1) + pad0 = (p + 1) // 2 + pad1 = p // 2 + else: + p = k.shape[0] - factor + pad0 = (p + 1) // 2 + pad1 = p // 2 + else: + k = resample_kernel + pad0, pad1 = 0, 0 + return k, pad0, pad1 + + +def upsample_2d(x, x_res, pad0, pad1, k, factor=2): + assert isinstance(factor, int) and factor >= 1 + return _simple_upfirdn_2d(x, x_res, k, up=factor, pad0=pad0, pad1=pad1) + + +def downsample_2d(x, x_res, pad0, pad1, k, factor=2): + assert isinstance(factor, int) and factor >= 1 + return _simple_upfirdn_2d(x, x_res, k, down=factor, pad0=pad0, pad1=pad1) + + +def upsample_conv_2d(x, x_res, w, convH, convW, pad0, pad1, k, factor=2): + assert isinstance(factor, int) and factor >= 1 + + # Check weight shape. + w = tf.convert_to_tensor(w) + assert w.shape.rank == 4 + # convH = w.shape[0] + # convW = w.shape[1] + inC = tf.shape(w)[2] + outC = tf.shape(w)[3] + assert convW == convH + + # Determine data dimensions. + stride = [1, 1, factor, factor] + output_shape = [tf.shape(x)[0], outC, (x_res - 1) * factor + convH, (int(x_res/2) - 1) * factor + convW] + num_groups = tf.shape(x)[1] // inC + + # Transpose weights. + w = tf.reshape(w, [convH, convW, inC, num_groups, -1]) + w = tf.transpose(w[::-1, ::-1], [0, 1, 4, 3, 2]) + w = tf.reshape(w, [convH, convW, -1, num_groups * inC]) + + # Execute. + x = tf.nn.conv2d_transpose(x, w, output_shape=output_shape, strides=stride, padding='VALID', data_format='NCHW') + new_x_res = output_shape[2] + return _simple_upfirdn_2d(x, new_x_res, k, pad0=pad0, pad1=pad1) + + +def conv_downsample_2d(x, x_res, w, convH, convW, pad0, pad1, k, factor=2): + assert isinstance(factor, int) and factor >= 1 + w = tf.convert_to_tensor(w) + # convH, convW, _inC, _outC = w.shape.as_list() + assert convW == convH + + s = [1, 1, factor, factor] + x = _simple_upfirdn_2d(x, x_res, k, pad0=pad0, pad1=pad1) + return tf.nn.conv2d(x, w, strides=s, padding='VALID', data_format='NCHW') + + +def _simple_upfirdn_2d(x, x_res, k, up=1, down=1, pad0=0, pad1=0): + assert x.shape.rank == 4 + y = x + y = tf.reshape(y, [-1, tf.shape(x)[2], tf.shape(x)[3], 1]) + y = upfirdn_2d_cuda(y, k, upx=up, upy=up, downx=down, downy=down, padx0=pad0, padx1=pad1, pady0=pad0, pady1=pad1) + y = tf.reshape(y, [-1, tf.shape(x)[1], tf.shape(y)[1], tf.shape(y)[2]]) + return y + + +def upfirdn_2d_cuda(x, k, upx, upy, downx, downy, padx0, padx1, pady0, pady1): + """Fast CUDA implementation of `upfirdn_2d()` using custom ops.""" + + x = tf.convert_to_tensor(x) + k = np.asarray(k, dtype=np.float32) + majorDim, inH, inW, minorDim = x.shape.as_list() + kernelH, kernelW = k.shape + assert inW >= 1 and inH >= 1 + assert kernelW >= 1 and kernelH >= 1 + assert isinstance(upx, int) and isinstance(upy, int) + assert isinstance(downx, int) and isinstance(downy, int) + assert isinstance(padx0, int) and isinstance(padx1, int) + assert isinstance(pady0, int) and isinstance(pady1, int) + + outW = (inW * upx + padx0 + padx1 - kernelW) // downx + 1 + outH = (inH * upy + pady0 + pady1 - kernelH) // downy + 1 + assert outW >= 1 and outH >= 1 + + kc = tf.constant(k, dtype=x.dtype) + gkc = tf.constant(k[::-1, ::-1], dtype=x.dtype) + gpadx0 = kernelW - padx0 - 1 + gpady0 = kernelH - pady0 - 1 + gpadx1 = inW * upx - outW * downx + padx0 - upx + 1 + gpady1 = inH * upy - outH * downy + pady0 - upy + 1 + + @tf.custom_gradient + def func(x): + y = _get_plugin().up_fir_dn2d(x=x, k=kc, upx=upx, upy=upy, downx=downx, downy=downy, padx0=padx0, padx1=padx1, pady0=pady0, pady1=pady1) + y.set_shape([majorDim, outH, outW, minorDim]) + @tf.custom_gradient + def grad(dy): + dx = _get_plugin().up_fir_dn2d(x=dy, k=gkc, upx=downx, upy=downy, downx=upx, downy=upy, padx0=gpadx0, padx1=gpadx1, pady0=gpady0, pady1=gpady1) + dx.set_shape([majorDim, inH, inW, minorDim]) + return dx, func + return y, grad + return func(x) diff --git a/stylegan2_ae/layers/dense.py b/stylegan2_ae/layers/dense.py new file mode 100644 index 0000000..912211d --- /dev/null +++ b/stylegan2_ae/layers/dense.py @@ -0,0 +1,37 @@ +import tensorflow as tf + +from stylegan2_ae.layers.commons import compute_runtime_coef + + +class Dense(tf.keras.layers.Layer): + def __init__(self, fmaps, gain, lrmul, **kwargs): + super(Dense, self).__init__(**kwargs) + self.fmaps = fmaps + self.gain = gain + self.lrmul = lrmul + + def build(self, input_shape): + fan_in = tf.reduce_prod(input_shape[1:]) + weight_shape = [fan_in, self.fmaps] + init_std, self.runtime_coef = compute_runtime_coef(weight_shape, self.gain, self.lrmul) + + w_init = tf.random.normal(shape=weight_shape, mean=0.0, stddev=init_std) + self.w = tf.Variable(w_init, name='w', trainable=True) + + def call(self, inputs, training=None, mask=None): + weight = self.runtime_coef * self.w + + c = tf.reduce_prod(tf.shape(inputs)[1:]) + x = tf.reshape(inputs, shape=[-1, c]) + x = tf.matmul(x, weight) + return x + + def get_config(self): + config = super(Dense, self).get_config() + config.update({ + 'fmaps': self.fmaps, + 'gain': self.gain, + 'lrmul': self.lrmul, + 'runtime_coef': self.runtime_coef, + }) + return config diff --git a/stylegan2_ae/layers/embedding.py b/stylegan2_ae/layers/embedding.py new file mode 100644 index 0000000..90e29c8 --- /dev/null +++ b/stylegan2_ae/layers/embedding.py @@ -0,0 +1,135 @@ +import tensorflow as tf + +from convnext.models.convnext_tf import ConvNeXt + +def rodrigues(r): + theta = tf.norm(r, axis=(1, 2), keepdims=True) + theta = tf.maximum(theta, 1e-7) + r_hat = r / theta + z_stick = tf.zeros(theta.get_shape().as_list()[0], dtype=tf.float32) + m = tf.stack( + (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick, + -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), axis=1) + m = tf.reshape(m, (-1, 3, 3)) + i_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros( + (theta.get_shape().as_list()[0], 3, 3), dtype=tf.float32) + # R = cos * i_cube + (1 - cos) * dot + tf.sin(theta) * m + R = i_cube + m * tf.sin(theta) + tf.matmul(m, m) * (1 - tf.cos(theta)) + return R + + +def run_rodrigues(pose): + batch_size = pose.shape[0] + R_out = rodrigues(tf.reshape(pose, (-1, 1, 3))) + R_out = tf.reshape(R_out, (batch_size, -1, 3, 3)) + return R_out + + +# if __name__ == '__main__': +# x = tf.ones((4, 72)) +# y = run_rodrigues(x) +# print() +# pass + + +class ParameterEmbedding(tf.keras.layers.Layer): + def __init__(self, embed_dim, **kwargs): + super(ParameterEmbedding, self).__init__(**kwargs) + self.embed_dim = embed_dim + + self.dense_1 = tf.keras.layers.Dense(128, activation='tanh', name='dense_1') + self.bn_1 = tf.keras.layers.BatchNormalization() + self.dense_2 = tf.keras.layers.Dense(512, activation='tanh', name='dense_2') + self.bn_2 = tf.keras.layers.BatchNormalization() + self.dense_3 = tf.keras.layers.Dense(self.embed_dim, activation='tanh', name='dense_3') + + def call(self, inputs, training=None, mask=None): + x = inputs + x = self.dense_1(x) + x = self.bn_1(x) + x = self.dense_2(x) + x = self.bn_2(x) + x = self.dense_3(x) + return x + + def get_config(self): + config = super(ParameterEmbedding, self).get_config() + config.update({ + 'embed_dim': self.embed_dim, + }) + return config + + +class PoseShapeEmbedding(tf.keras.layers.Layer): + def __init__(self, **kwargs): + super(PoseShapeEmbedding, self).__init__(**kwargs) + + # self.pose_embedding = ParameterEmbedding(int(self.embed_dim/2)) + # self.shape_embedding = ParameterEmbedding(int(self.embed_dim/2)) + self.flatten = tf.keras.layers.Flatten() + + def call(self, inputs, training=None, mask=None): + pose, shape = inputs + rotations = run_rodrigues(pose) + rotations = self.flatten(rotations) + # pose_latent = self.pose_embedding(rotations, training) + # shape_latent = self.shape_embedding(shape, training) + return tf.concat((rotations, shape), axis=-1) + + def get_config(self): + config = super(PoseShapeEmbedding, self).get_config() + config.update({ + 'embed_dim': 216 + 11, + }) + return config + + +class GarmentEmbedding(tf.keras.layers.Layer): + def __init__(self, embed_dim, **kwargs): + super(GarmentEmbedding, self).__init__(**kwargs) + self.embed_dim = embed_dim + + # TODO: channels_first? + + self.encoder = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 256, self.embed_dim], include_top=False) + # self.encoder(tf.keras.Input(shape=(*input_shape, 3))) + self.pooling = tf.keras.layers.GlobalAveragePooling2D() + + def call(self, inputs, training=None, mask=None): + x = inputs + + x = tf.transpose(x, [0,2,3,1]) + x = self.encoder(x, training=training) + x = self.pooling(x) + + return x + + def get_config(self): + config = super(GarmentEmbedding, self).get_config() + config.update({ + 'embed_dim': self.embed_dim, + }) + return config + + +# class LabelEmbedding(tf.keras.layers.Layer): +# def __init__(self, embed_dim, **kwargs): +# super(LabelEmbedding, self).__init__(**kwargs) +# self.embed_dim = embed_dim + +# def build(self, input_shape): +# weight_shape = [input_shape[1], self.embed_dim] +# # tf 1.15 mean(0.0), std(1.0) default value of tf.initializers.random_normal() +# w_init = tf.random.normal(shape=weight_shape, mean=0.0, stddev=1.0) +# self.w = tf.Variable(w_init, name='w', trainable=True) + +# def call(self, inputs, training=None, mask=None): +# x = tf.matmul(inputs, self.w) +# return x + +# def get_config(self): +# config = super(LabelEmbedding, self).get_config() +# config.update({ +# 'embed_dim': self.embed_dim, +# }) +# return config diff --git a/stylegan2_ae/layers/from_rgb.py b/stylegan2_ae/layers/from_rgb.py new file mode 100644 index 0000000..b69e677 --- /dev/null +++ b/stylegan2_ae/layers/from_rgb.py @@ -0,0 +1,28 @@ +import tensorflow as tf + +from stylegan2_ae.layers.conv import Conv2D +from stylegan2_ae.layers.bias_act import BiasAct + + +class FromRGB(tf.keras.layers.Layer): + def __init__(self, fmaps, res, **kwargs): + super(FromRGB, self).__init__(**kwargs) + self.fmaps = fmaps + self.res = res + + self.conv = Conv2D(in_res=res, in_fmaps=3, fmaps=self.fmaps, kernel=1, up=False, down=False, + resample_kernel=None, gain=1.0, lrmul=1.0, name='conv') + self.apply_bias_act = BiasAct(lrmul=1.0, act='lrelu', name='bias') + + def call(self, inputs, training=None, mask=None): + y = self.conv(inputs) + y = self.apply_bias_act(y) + return y + + def get_config(self): + config = super(FromRGB, self).get_config() + config.update({ + 'fmaps': self.fmaps, + 'res': self.res, + }) + return config diff --git a/stylegan2_ae/layers/mapping_block.py b/stylegan2_ae/layers/mapping_block.py new file mode 100644 index 0000000..eb6bbfd --- /dev/null +++ b/stylegan2_ae/layers/mapping_block.py @@ -0,0 +1,48 @@ +import tensorflow as tf + +from stylegan2_ae.layers.embedding import PoseShapeEmbedding, GarmentEmbedding +from stylegan2_ae.layers.dense import Dense +from stylegan2_ae.layers.bias_act import BiasAct + + +class Mapping(tf.keras.layers.Layer): + def __init__(self, w_dim, labels_dim, n_mapping, name, **kwargs): + super(Mapping, self).__init__(name=name, **kwargs) + self.w_dim = w_dim + self.labels_dim = labels_dim + self.n_mapping = n_mapping + self.gain = 1.0 + self.lrmul = 0.01 + + self.normalize = tf.keras.layers.Lambda(lambda x: x * tf.math.rsqrt(tf.reduce_mean(tf.square(x), axis=1, keepdims=True) + 1e-8)) + + self.dense_layers = list() + self.bias_act_layers = list() + for ii in range(self.n_mapping): + self.dense_layers.append(Dense(w_dim, gain=self.gain, lrmul=self.lrmul, name='dense_{:d}'.format(ii))) + self.bias_act_layers.append(BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_{:d}'.format(ii))) + + def call(self, inputs, training=None, mask=None): + latents = inputs + x = latents + + # normalize inputs + x = self.normalize(x) + + # apply mapping blocks + for dense, apply_bias_act in zip(self.dense_layers, self.bias_act_layers): + x = dense(x) + x = apply_bias_act(x) + + return x + + def get_config(self): + config = super(Mapping, self).get_config() + config.update({ + 'w_dim': self.w_dim, + 'labels_dim': self.labels_dim, + 'n_mapping': self.n_mapping, + 'gain': self.gain, + 'lrmul': self.lrmul, + }) + return config diff --git a/stylegan2_ae/layers/mini_batch_std.py b/stylegan2_ae/layers/mini_batch_std.py new file mode 100644 index 0000000..30389ed --- /dev/null +++ b/stylegan2_ae/layers/mini_batch_std.py @@ -0,0 +1,33 @@ +import tensorflow as tf + + +class MinibatchStd(tf.keras.layers.Layer): + def __init__(self, group_size, num_new_features, **kwargs): + super(MinibatchStd, self).__init__(**kwargs) + self.group_size = group_size + self.num_new_features = num_new_features + + def call(self, inputs, training=None, mask=None): + s = tf.shape(inputs) + group_size = tf.minimum(self.group_size, s[0]) + + y = tf.reshape(inputs, [group_size, -1, self.num_new_features, s[1] // self.num_new_features, s[2], s[3]]) + y = tf.cast(y, tf.float32) + y -= tf.reduce_mean(y, axis=0, keepdims=True) + y = tf.reduce_mean(tf.square(y), axis=0) + y = tf.sqrt(y + 1e-8) + y = tf.reduce_mean(y, axis=[2, 3, 4], keepdims=True) + y = tf.reduce_mean(y, axis=[2]) + y = tf.cast(y, inputs.dtype) + y = tf.tile(y, [group_size, 1, s[2], s[3]]) + + x = tf.concat([inputs, y], axis=1) + return x + + def get_config(self): + config = super(MinibatchStd, self).get_config() + config.update({ + 'group_size': self.group_size, + 'num_new_features': self.num_new_features, + }) + return config diff --git a/stylegan2_ae/layers/modulated_conv2d.py b/stylegan2_ae/layers/modulated_conv2d.py new file mode 100644 index 0000000..6d37ea7 --- /dev/null +++ b/stylegan2_ae/layers/modulated_conv2d.py @@ -0,0 +1,105 @@ +import tensorflow as tf + +from stylegan2_ae.layers.commons import compute_runtime_coef +from stylegan2_ae.layers.dense import Dense +from stylegan2_ae.layers.bias_act import BiasAct +from stylegan2_ae.layers.cuda.upfirdn_2d_v2 import upsample_conv_2d, conv_downsample_2d, compute_paddings + + +class ModulatedConv2D(tf.keras.layers.Layer): + def __init__(self, in_res, in_fmaps, fmaps, kernel, up, down, demodulate, resample_kernel, gain, lrmul, + fused_modconv, **kwargs): + super(ModulatedConv2D, self).__init__(**kwargs) + assert not (up and down) + + self.in_res = in_res + self.in_fmaps = in_fmaps + self.fmaps = fmaps + self.kernel = kernel + self.demodulate = demodulate + self.up = up + self.down = down + self.fused_modconv = fused_modconv + self.gain = gain + self.lrmul = lrmul + # self.resample_kernel = resample_kernel + + self.k, self.pad0, self.pad1 = compute_paddings(resample_kernel, self.kernel, up, down, is_conv=True) + + # self.factor = 2 + self.mod_dense = Dense(self.in_fmaps, gain=1.0, lrmul=1.0, name='mod_dense') + self.mod_bias = BiasAct(lrmul=1.0, act='linear', name='mod_bias') + + def build(self, input_shape): + # x_shape, w_shape = input_shape[0], input_shape[1] + weight_shape = [self.kernel, self.kernel, self.in_fmaps, self.fmaps] + init_std, self.runtime_coef = compute_runtime_coef(weight_shape, self.gain, self.lrmul) + + # [kkIO] + w_init = tf.random.normal(shape=weight_shape, mean=0.0, stddev=init_std) + self.w = tf.Variable(w_init, name='w', trainable=True) + + def call(self, inputs, training=None, mask=None): + x, y = inputs + # height, width = tf.shape(x)[2], tf.shape(x)[3] + + # prepare weights: [BkkIO] Introduce minibatch dimension + w = self.runtime_coef * self.w + ww = w[tf.newaxis] + + # Modulate + s = self.mod_dense(y) # [BI] + s = self.mod_bias(s) + 1.0 # [BI] + ww *= s[:, tf.newaxis, tf.newaxis, :, tf.newaxis] # [BkkIO] + + if self.demodulate: + d = tf.math.rsqrt(tf.reduce_sum(tf.square(ww), axis=[1, 2, 3]) + 1e-8) # [BO] + ww *= d[:, tf.newaxis, tf.newaxis, tf.newaxis, :] # [BkkIO] + + if self.fused_modconv: + # Fused => reshape minibatch to convolution groups + x_shape = tf.shape(x) + ww_shape = tf.shape(ww) + x = tf.reshape(x, [1, -1, x_shape[2], x_shape[3]]) + w = tf.reshape(tf.transpose(ww, [1, 2, 3, 0, 4]), [ww_shape[1], ww_shape[2], ww_shape[3], -1]) + else: + # [BIhw] Not fused => scale input activations + x *= s[:, :, tf.newaxis, tf.newaxis] + + # Convolution with optional up/downsampling. + if self.up: + x = upsample_conv_2d(x, self.in_res, w, self.kernel, self.kernel, self.pad0, self.pad1, self.k) + elif self.down: + x = conv_downsample_2d(x, self.in_res, w, self.kernel, self.kernel, self.pad0, self.pad1, self.k) + else: + x = tf.nn.conv2d(x, w, data_format='NCHW', strides=[1, 1, 1, 1], padding='SAME') + + # Reshape/scale output + if self.fused_modconv: + # Fused => reshape convolution groups back to minibatch + x_shape = tf.shape(x) + x = tf.reshape(x, [-1, self.fmaps, x_shape[2], x_shape[3]]) + elif self.demodulate: + # [BOhw] Not fused => scale output activations + x *= d[:, :, tf.newaxis, tf.newaxis] + return x + + def get_config(self): + config = super(ModulatedConv2D, self).get_config() + config.update({ + 'in_res': self.in_res, + 'in_fmaps': self.in_fmaps, + 'fmaps': self.fmaps, + 'kernel': self.kernel, + 'demodulate': self.demodulate, + 'fused_modconv': self.fused_modconv, + 'gain': self.gain, + 'lrmul': self.lrmul, + 'up': self.up, + 'down': self.down, + 'k': self.k, + 'pad0': self.pad0, + 'pad1': self.pad1, + 'runtime_coef': self.runtime_coef, + }) + return config diff --git a/stylegan2_ae/layers/noise.py b/stylegan2_ae/layers/noise.py new file mode 100644 index 0000000..3255073 --- /dev/null +++ b/stylegan2_ae/layers/noise.py @@ -0,0 +1,25 @@ +import tensorflow as tf + + +class Noise(tf.keras.layers.Layer): + def __init__(self, **kwargs): + super(Noise, self).__init__(**kwargs) + + def build(self, input_shape): + w_init = tf.zeros(shape=(), dtype=tf.dtypes.float32) + self.noise_strength = tf.Variable(w_init, trainable=True, name='w') + + def call(self, inputs, noise=None, training=None, mask=None): + x_shape = tf.shape(inputs) + + # noise: [1, 1, x_shape[2], x_shape[3]] or None + if noise is None: + noise = tf.random.normal(shape=(x_shape[0], 1, x_shape[2], x_shape[3]), dtype=tf.dtypes.float32) + + x = inputs + noise * self.noise_strength + return x + + def get_config(self): + config = super(Noise, self).get_config() + config.update({}) + return config diff --git a/stylegan2_ae/layers/synthesis_block.py b/stylegan2_ae/layers/synthesis_block.py new file mode 100644 index 0000000..f7b61a4 --- /dev/null +++ b/stylegan2_ae/layers/synthesis_block.py @@ -0,0 +1,162 @@ +import tensorflow as tf + +from stylegan2_ae.layers.modulated_conv2d import ModulatedConv2D +from stylegan2_ae.layers.bias_act import BiasAct +from stylegan2_ae.layers.noise import Noise +from stylegan2_ae.layers.to_rgb import ToRGB +from stylegan2_ae.layers.cuda.upfirdn_2d_v2 import upsample_2d, compute_paddings + + +class SynthesisConstBlock(tf.keras.layers.Layer): + def __init__(self, fmaps, res, **kwargs): + super(SynthesisConstBlock, self).__init__(**kwargs) + assert res == 4 + self.res = res + self.fmaps = fmaps + self.gain = 1.0 + self.lrmul = 1.0 + + # conv block + self.conv = ModulatedConv2D(in_res=res, in_fmaps=self.fmaps, fmaps=self.fmaps, kernel=3, up=False, down=False, + demodulate=True, resample_kernel=[1, 3, 3, 1], gain=self.gain, lrmul=self.lrmul, + fused_modconv=True, name='conv') + self.apply_noise = Noise(name='noise') + self.apply_bias_act = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias') + + def build(self, input_shape): + # starting const variable + # tf 1.15 mean(0.0), std(1.0) default value of tf.initializers.random_normal() + const_init = tf.random.normal(shape=(1, self.fmaps, self.res, int(self.res/2)), mean=0.0, stddev=1.0) + self.const = tf.Variable(const_init, name='const', trainable=True) + + def call(self, inputs, training=None, mask=None): + w0 = inputs + batch_size = tf.shape(w0)[0] + + # const block + x = tf.tile(self.const, [batch_size, 1, 1, 1]) + + # conv block + x = self.conv([x, w0]) + x = self.apply_noise(x) + x = self.apply_bias_act(x) + return x + + def get_config(self): + config = super(SynthesisConstBlock, self).get_config() + config.update({ + 'res': self.res, + 'fmaps': self.fmaps, + 'gain': self.gain, + 'lrmul': self.lrmul, + }) + return config + + +class SynthesisBlock(tf.keras.layers.Layer): + def __init__(self, in_ch, fmaps, res, **kwargs): + super(SynthesisBlock, self).__init__(**kwargs) + self.in_ch = in_ch + self.fmaps = fmaps + self.res = res + self.gain = 1.0 + self.lrmul = 1.0 + + # conv0 up + self.conv_0 = ModulatedConv2D(in_res=res//2, in_fmaps=self.in_ch, fmaps=self.fmaps, kernel=3, up=True, down=False, + demodulate=True, resample_kernel=[1, 3, 3, 1], gain=self.gain, lrmul=self.lrmul, + fused_modconv=True, name='conv_0') + self.apply_noise_0 = Noise(name='noise_0') + self.apply_bias_act_0 = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_0') + + # conv block + self.conv_1 = ModulatedConv2D(in_res=res, in_fmaps=self.fmaps, fmaps=self.fmaps, kernel=3, up=False, down=False, + demodulate=True, resample_kernel=[1, 3, 3, 1], gain=self.gain, lrmul=self.lrmul, + fused_modconv=True, name='conv_1') + self.apply_noise_1 = Noise(name='noise_1') + self.apply_bias_act_1 = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_1') + + def call(self, inputs, training=None, mask=None): + x, w0, w1 = inputs + + # conv0 up + x = self.conv_0([x, w0]) + x = self.apply_noise_0(x) + x = self.apply_bias_act_0(x) + + # conv block + x = self.conv_1([x, w1]) + x = self.apply_noise_1(x) + x = self.apply_bias_act_1(x) + return x + + def get_config(self): + config = super(SynthesisBlock, self).get_config() + config.update({ + 'in_ch': self.in_ch, + 'res': self.res, + 'fmaps': self.fmaps, + 'gain': self.gain, + 'lrmul': self.lrmul, + }) + return config + + +class Synthesis(tf.keras.layers.Layer): + def __init__(self, resolutions, featuremaps, name, **kwargs): + super(Synthesis, self).__init__(name=name, **kwargs) + self.resolutions = resolutions + self.featuremaps = featuremaps + # self.resample_kernel = [1, 3, 3, 1] + + self.k, self.pad0, self.pad1 = compute_paddings([1, 3, 3, 1], None, up=True, down=False, is_conv=False) + + # initial layer + res, n_f = resolutions[0], featuremaps[0] + self.initial_block = SynthesisConstBlock(fmaps=n_f, res=res, name='{:d}x{:d}/const'.format(res, res)) + self.initial_torgb = ToRGB(in_ch=n_f, res=res, name='{:d}x{:d}/ToRGB'.format(res, res)) + + # stack generator block with lerp block + prev_n_f = n_f + self.blocks = list() + self.torgbs = list() + for res, n_f in zip(self.resolutions[1:], self.featuremaps[1:]): + self.blocks.append(SynthesisBlock(in_ch=prev_n_f, fmaps=n_f, res=res, + name='{:d}x{:d}/block'.format(res, res))) + self.torgbs.append(ToRGB(in_ch=n_f, res=res * 2, name='{:d}x{:d}/ToRGB'.format(res, res))) + prev_n_f = n_f + + def call(self, inputs, training=None, mask=None): + w_broadcasted = inputs + + # initial layer + w0, w1 = w_broadcasted[:, 0], w_broadcasted[:, 1] + x = self.initial_block(w0) + y = self.initial_torgb([x, w1]) + + layer_index = 1 + for block, torgb in zip(self.blocks, self.torgbs): + w0 = w_broadcasted[:, layer_index] + w1 = w_broadcasted[:, layer_index + 1] + w2 = w_broadcasted[:, layer_index + 2] + + y_res = block.res // 2 + x = block([x, w0, w1]) + y = upsample_2d(y, y_res, self.pad0, self.pad1, self.k) + y = y + torgb([x, w2]) + + layer_index += 2 + + images_out = y + return images_out + + def get_config(self): + config = super(Synthesis, self).get_config() + config.update({ + 'resolutions': self.resolutions, + 'featuremaps': self.featuremaps, + 'k': self.k, + 'pad0': self.pad0, + 'pad1': self.pad1, + }) + return config diff --git a/stylegan2_ae/layers/to_rgb.py b/stylegan2_ae/layers/to_rgb.py new file mode 100644 index 0000000..78a2121 --- /dev/null +++ b/stylegan2_ae/layers/to_rgb.py @@ -0,0 +1,30 @@ +import tensorflow as tf + +from stylegan2.layers.modulated_conv2d import ModulatedConv2D +from stylegan2.layers.bias_act import BiasAct + + +class ToRGB(tf.keras.layers.Layer): + def __init__(self, in_ch, res, **kwargs): + super(ToRGB, self).__init__(**kwargs) + self.in_ch = in_ch + self.res = res + + self.conv = ModulatedConv2D(in_res=res, in_fmaps=in_ch, fmaps=3, kernel=1, up=False, down=False, demodulate=False, + resample_kernel=None, gain=1.0, lrmul=1.0, fused_modconv=True, name='conv') + self.apply_bias = BiasAct(lrmul=1.0, act='linear', name='bias') + + def call(self, inputs, training=None, mask=None): + x, w = inputs + + x = self.conv([x, w]) + x = self.apply_bias(x) + return x + + def get_config(self): + config = super(ToRGB, self).get_config() + config.update({ + 'in_ch': self.in_ch, + 'res': self.res, + }) + return config diff --git a/stylegan2_ae/utils.py b/stylegan2_ae/utils.py new file mode 100644 index 0000000..bab4172 --- /dev/null +++ b/stylegan2_ae/utils.py @@ -0,0 +1,112 @@ +from typing import Any +import numpy as np +import tensorflow as tf + + +class EasyDict(dict): + """Convenience class that behaves like a dict but allows access with the attribute syntax.""" + + def __getattr__(self, name: str) -> Any: + try: + return self[name] + except KeyError: + raise AttributeError(name) + + def __setattr__(self, name: str, value: Any) -> None: + self[name] = value + + def __delattr__(self, name: str) -> None: + del self[name] + + +def lerp(a, b, t): + out = a + (b - a) * t + return out + + +def lerp_clip(a, b, t): + out = a + (b - a) * tf.clip_by_value(t, 0.0, 1.0) + return out + + +def adjust_dynamic_range(images, range_in, range_out, out_dtype): + scale = (range_out[1] - range_out[0]) / (range_in[1] - range_in[0]) + bias = range_out[0] - range_in[0] * scale + images = images * scale + bias + images = tf.clip_by_value(images, range_out[0], range_out[1]) + images = tf.cast(images, dtype=out_dtype) + return images + + +def random_flip_left_right_nchw(images): + s = tf.shape(images) + mask = tf.random.uniform([s[0], 1, 1, 1], 0.0, 1.0) + mask = tf.tile(mask, [1, s[1], s[2], s[3]]) + images = tf.where(mask < 0.5, images, tf.reverse(images, axis=[3])) + return images + + +def preprocess_fit_train_image(images, res): + images = adjust_dynamic_range(images, range_in=(0.0, 255.0), range_out=(-1.0, 1.0), out_dtype=tf.dtypes.float32) + images = random_flip_left_right_nchw(images) + images.set_shape([None, 3, res, res]) + return images + + +def postprocess_images(images): + images = adjust_dynamic_range(images, range_in=(-1.0, 1.0), range_out=(0.0, 255.0), out_dtype=tf.dtypes.float32) + images = tf.transpose(images, [0, 2, 3, 1]) + images = tf.cast(images, dtype=tf.dtypes.uint8) + return images + + +def merge_batch_images(images, res, rows, cols): + batch_size = images.shape[0] + assert rows * cols == batch_size + canvas = np.zeros(shape=[res * rows, res * cols, 3], dtype=np.uint8) + for row in range(rows): + y_start = row * res + for col in range(cols): + x_start = col * res + index = col + row * cols + canvas[y_start:y_start + res, x_start:x_start + res, :] = images[index, :, :, :] + return canvas + + +def main(): + # out_dtype = tf.dtypes.float32 + # range_in = (0.0, 255.0) + # range_out = (-1.0, 1.0) + # # images = np.ones((2, 2), dtype=np.float32) * 255.0 + # images1 = np.ones((2, 2), dtype=np.float32) * 0.0 + # images2 = np.ones((2, 2), dtype=np.float32) * 127.5 + # images3 = np.ones((2, 2), dtype=np.float32) * 255.0 + # adjusted1 = adjust_dynamic_range(images1, range_in, range_out, out_dtype) + # adjusted2 = adjust_dynamic_range(images2, range_in, range_out, out_dtype) + # adjusted3 = adjust_dynamic_range(images3, range_in, range_out, out_dtype) + # print(adjusted1) + # print(adjusted2) + # print(adjusted3) + # + # out_dtype = tf.dtypes.uint8 + # range_in = (-1.0, 1.0) + # range_out = (0.0, 255.0) + # images1 = np.ones((2, 2), dtype=np.float32) * -1.0 + # images2 = np.ones((2, 2), dtype=np.float32) * 0.0 + # images3 = np.ones((2, 2), dtype=np.float32) * 1.0 + # adjusted1 = adjust_dynamic_range(images1, range_in, range_out, out_dtype) + # adjusted2 = adjust_dynamic_range(images2, range_in, range_out, out_dtype) + # adjusted3 = adjust_dynamic_range(images3, range_in, range_out, out_dtype) + # print(adjusted1) + # print(adjusted2) + # print(adjusted3) + + batch_size = 8 + res = 128 + fake_images = np.ones(shape=(batch_size, res, res, 3), dtype=np.uint8) + merge_batch_images(fake_images, res, rows=4, cols=2) + return + + +if __name__ == '__main__': + main()