add ae sub folder
This commit is contained in:
@@ -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)
|
||||
@@ -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]
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
@@ -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 <stdio.h>
|
||||
|
||||
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 <class T>
|
||||
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 <class T>
|
||||
static __global__ void FusedBiasActKernel(const FusedBiasActKernelParams<T> 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 <class T>
|
||||
struct FusedBiasActOp : public OpKernel
|
||||
{
|
||||
FusedBiasActKernelParams<T> 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<T> p = m_attribs;
|
||||
cudaStream_t stream = ctx->eigen_device<Eigen::GpuDevice>().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<T>().data();
|
||||
p.b = (b.NumElements()) ? b.flat<T>().data() : NULL;
|
||||
p.ref = (ref.NumElements()) ? ref.flat<T>().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<T>().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<T>, 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<float>("T"), FusedBiasActOp<float>);
|
||||
REGISTER_KERNEL_BUILDER(Name("FusedBiasAct").Device(DEVICE_GPU).TypeConstraint<Eigen::half>("T"), FusedBiasActOp<Eigen::half>);
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
@@ -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)
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
@@ -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 <stdio.h>
|
||||
|
||||
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 <class T>
|
||||
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 <class T>
|
||||
static __global__ void UpFirDn2DKernel_large(const UpFirDn2DKernelParams<T> 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 <class T, int upx, int upy, int downx, int downy, int kernelW, int kernelH, int tileOutW, int tileOutH>
|
||||
static __global__ void UpFirDn2DKernel_small(const UpFirDn2DKernelParams<T> 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 <class T>
|
||||
struct UpFirDn2DOp : public OpKernel
|
||||
{
|
||||
UpFirDn2DKernelParams<T> 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<T> p = m_attribs;
|
||||
cudaStream_t stream = ctx->eigen_device<Eigen::GpuDevice>().stream();
|
||||
|
||||
const Tensor& x = ctx->input(0); // [majorDim, inH, inW, minorDim]
|
||||
const Tensor& k = ctx->input(1); // [kernelH, kernelW]
|
||||
p.x = x.flat<T>().data();
|
||||
p.k = k.flat<T>().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<T>().data();
|
||||
OP_REQUIRES(ctx, y->NumElements() <= kint32max, errors::InvalidArgument("output too large"));
|
||||
|
||||
// Choose CUDA kernel to use.
|
||||
void* cudaKernel = (void*)UpFirDn2DKernel_large<T>;
|
||||
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<T, 1,1, 1,1, 7,7, 64,16>; 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<T, 1,1, 1,1, 6,6, 64,16>; 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<T, 1,1, 1,1, 5,5, 64,16>; 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<T, 1,1, 1,1, 4,4, 64,16>; 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<T, 1,1, 1,1, 3,3, 64,16>; 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<T, 2,2, 1,1, 8,8, 64,16>; 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<T, 2,2, 1,1, 6,6, 64,16>; 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<T, 2,2, 1,1, 4,4, 64,16>; 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<T, 2,2, 1,1, 2,2, 64,16>; 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<T, 1,1, 2,2, 8,8, 32,8>; 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<T, 1,1, 2,2, 6,6, 32,8>; 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<T, 1,1, 2,2, 4,4, 32,8>; 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<T, 1,1, 2,2, 2,2, 32,8>; 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<float>("T"), UpFirDn2DOp<float>);
|
||||
REGISTER_KERNEL_BUILDER(Name("UpFirDn2D").Device(DEVICE_GPU).TypeConstraint<Eigen::half>("T"), UpFirDn2DOp<Eigen::half>);
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
@@ -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
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user