init
This commit is contained in:
@@ -0,0 +1,277 @@
|
||||
import os
|
||||
import tensorflow as tf
|
||||
|
||||
from load_models import load_generator, load_discriminator
|
||||
|
||||
|
||||
def handle_mapping(w_name, is_g_clone):
|
||||
def extract_info(name):
|
||||
splitted = name.split('/')
|
||||
index = splitted.index('g_mapping')
|
||||
indicator = splitted[index + 1]
|
||||
val = indicator.split('_')[-1]
|
||||
return val
|
||||
|
||||
level = extract_info(w_name)
|
||||
o_prefix = f'G_mapping_1' if is_g_clone else f'G_mapping'
|
||||
if 'w' in w_name:
|
||||
official_var_name = f'{o_prefix}/Dense{level}/weight'
|
||||
else:
|
||||
official_var_name = f'{o_prefix}/Dense{level}/bias'
|
||||
return official_var_name
|
||||
|
||||
|
||||
def to_rgb_layer(name, r, o_prefix):
|
||||
if 'conv/w:0' in name:
|
||||
o_name = f'{o_prefix}/{r}x{r}/ToRGB/weight'
|
||||
elif 'mod_dense/w:0' in name:
|
||||
o_name = f'{o_prefix}/{r}x{r}/ToRGB/mod_weight'
|
||||
elif 'mod_bias/b:0' in name:
|
||||
o_name = f'{o_prefix}/{r}x{r}/ToRGB/mod_bias'
|
||||
else: # if 'bias/b:0' in name:
|
||||
o_name = f'{o_prefix}/{r}x{r}/ToRGB/bias'
|
||||
return o_name
|
||||
|
||||
|
||||
def handle_block_layer(name, r, o_prefix):
|
||||
if 'conv_0/w:0' in name:
|
||||
o_name = f'{o_prefix}/{r}x{r}/Conv0_up/weight'
|
||||
elif 'conv_0/mod_dense/w:0' in name:
|
||||
o_name = f'{o_prefix}/{r}x{r}/Conv0_up/mod_weight'
|
||||
elif 'conv_0/mod_bias/b:0' in name:
|
||||
o_name = f'{o_prefix}/{r}x{r}/Conv0_up/mod_bias'
|
||||
elif 'noise_0/w:0' in name:
|
||||
o_name = f'{o_prefix}/{r}x{r}/Conv0_up/noise_strength'
|
||||
elif 'bias_0/b:0' in name:
|
||||
o_name = f'{o_prefix}/{r}x{r}/Conv0_up/bias'
|
||||
elif 'conv_1/w:0' in name:
|
||||
o_name = f'{o_prefix}/{r}x{r}/Conv1/weight'
|
||||
elif 'conv_1/mod_dense/w:0' in name:
|
||||
o_name = f'{o_prefix}/{r}x{r}/Conv1/mod_weight'
|
||||
elif 'conv_1/mod_bias/b:0' in name:
|
||||
o_name = f'{o_prefix}/{r}x{r}/Conv1/mod_bias'
|
||||
elif 'noise_1/w:0' in name:
|
||||
o_name = f'{o_prefix}/{r}x{r}/Conv1/noise_strength'
|
||||
else: # if 'bias_1/b:0' in name:
|
||||
o_name = f'{o_prefix}/{r}x{r}/Conv1/bias'
|
||||
return o_name
|
||||
|
||||
|
||||
def handle_const_layer(name, o_prefix):
|
||||
if 'const:0' in name:
|
||||
o_name = f'{o_prefix}/4x4/Const/const'
|
||||
elif 'conv/w:0' in name:
|
||||
o_name = f'{o_prefix}/4x4/Conv/weight'
|
||||
elif 'mod_dense/w:0' in name:
|
||||
o_name = f'{o_prefix}/4x4/Conv/mod_weight'
|
||||
elif 'mod_bias/b:0' in name:
|
||||
o_name = f'{o_prefix}/4x4/Conv/mod_bias'
|
||||
elif 'noise/w:0' in name:
|
||||
o_name = f'{o_prefix}/4x4/Conv/noise_strength'
|
||||
else:
|
||||
o_name = f'{o_prefix}/4x4/Conv/bias'
|
||||
return o_name
|
||||
|
||||
|
||||
def handle_synthesis(w_name, is_g_clone):
|
||||
def extract_info(name):
|
||||
splitted = name.split('/')
|
||||
index = splitted.index('g_synthesis')
|
||||
indicator1 = splitted[index + 1]
|
||||
indicator2 = splitted[index + 2]
|
||||
r = indicator1.split('x')[1]
|
||||
d = indicator2
|
||||
return r, d
|
||||
|
||||
res, divider = extract_info(w_name)
|
||||
o_prefix = f'G_synthesis_1' if is_g_clone else f'G_synthesis'
|
||||
if divider == 'ToRGB':
|
||||
official_var_name = to_rgb_layer(w_name, res, o_prefix)
|
||||
elif divider == 'block':
|
||||
official_var_name = handle_block_layer(w_name, res, o_prefix)
|
||||
else: # const
|
||||
official_var_name = handle_const_layer(w_name, o_prefix)
|
||||
return official_var_name
|
||||
|
||||
|
||||
def handle_discriminator_layer(w_name):
|
||||
def extract_info(name):
|
||||
splitted = name.split('/')
|
||||
resolution = splitted[1]
|
||||
resolution = resolution.split('x')[0]
|
||||
return resolution
|
||||
|
||||
res = extract_info(w_name)
|
||||
if 'last_dense' in w_name:
|
||||
o_name = 'D/Output/weight'
|
||||
elif 'last_bias' in w_name:
|
||||
o_name = 'D/Output/bias'
|
||||
elif 'FromRGB' in w_name:
|
||||
o_name = 'D/1024x1024/FromRGB/weight' if 'conv_' in w_name else 'D/1024x1024/FromRGB/bias'
|
||||
elif 'skip' in w_name:
|
||||
o_name = f'D/{res}x{res}/Skip/weight'
|
||||
elif 'dense_' in w_name:
|
||||
o_name = f'D/4x4/Dense0/weight'
|
||||
elif 'conv_' in w_name:
|
||||
if res == '4':
|
||||
o_name = f'D/{res}x{res}/Conv/weight'
|
||||
else:
|
||||
o_name = f'D/{res}x{res}/Conv0/weight' if '_0' in w_name else f'D/{res}x{res}/Conv1_down/weight'
|
||||
elif 'bias_' in w_name:
|
||||
if res == '4':
|
||||
o_name = f'D/{res}x{res}/Conv/bias' if '_0' in w_name else f'D/{res}x{res}/Dense0/bias'
|
||||
else:
|
||||
o_name = f'D/{res}x{res}/Conv0/bias' if '_0' in w_name else f'D/{res}x{res}/Conv1_down/bias'
|
||||
else:
|
||||
raise ValueError('Something went wrong!!')
|
||||
return o_name
|
||||
|
||||
|
||||
def variable_name_mapper_g(g, is_g_clone):
|
||||
name_mapper = dict()
|
||||
for w in g.weights:
|
||||
w_name, w_shape = w.name, w.shape
|
||||
|
||||
# mapping layer
|
||||
if 'g_mapping' in w_name:
|
||||
official_var_name = handle_mapping(w_name, is_g_clone)
|
||||
elif 'g_synthesis' in w_name:
|
||||
official_var_name = handle_synthesis(w_name, is_g_clone)
|
||||
else:
|
||||
# w_avg
|
||||
official_var_name = 'Gs/dlatent_avg' if is_g_clone else 'G/dlatent_avg'
|
||||
|
||||
name_mapper[official_var_name] = w
|
||||
return name_mapper
|
||||
|
||||
|
||||
def variable_name_mapper_d(d):
|
||||
name_mapper = dict()
|
||||
for w in d.weights:
|
||||
w_name, w_shape = w.name, w.shape
|
||||
|
||||
official_var_name = handle_discriminator_layer(w_name)
|
||||
name_mapper[official_var_name] = w
|
||||
return name_mapper
|
||||
|
||||
|
||||
def check_shape(name_mapper, official_vars):
|
||||
for official_name, v in name_mapper.items():
|
||||
official_shape = [s for n, s in official_vars if n == official_name][0]
|
||||
|
||||
if official_shape == v.shape:
|
||||
print('{}: shape matches'.format(official_name))
|
||||
else:
|
||||
# print(f'Official: {official_name} -> {official_shape}')
|
||||
# print(f'Current: {v.name} -> {v.shape}')
|
||||
raise ValueError('{}: wrong shape'.format(official_name))
|
||||
return
|
||||
|
||||
|
||||
def convert_official_generator_weights(ckpt_dir, is_g_clone, use_custom_cuda):
|
||||
generator = load_generator(g_params=None, is_g_clone=is_g_clone, ckpt_dir=None, custom_cuda=use_custom_cuda)
|
||||
|
||||
# restore official ones to current implementation
|
||||
official_checkpoint = tf.train.latest_checkpoint('./official-pretrained')
|
||||
official_vars = tf.train.list_variables(official_checkpoint)
|
||||
|
||||
# get name mapper
|
||||
name_mapper = variable_name_mapper_g(generator, is_g_clone=is_g_clone)
|
||||
for name_g, tvar in name_mapper.items():
|
||||
print(f'{name_g}: {tvar.name}')
|
||||
|
||||
# check shape
|
||||
check_shape(name_mapper, official_vars)
|
||||
|
||||
# restore
|
||||
tf.compat.v1.train.init_from_checkpoint(official_checkpoint, assignment_map=name_mapper)
|
||||
|
||||
# save
|
||||
if is_g_clone:
|
||||
ckpt = tf.train.Checkpoint(g_clone=generator)
|
||||
else:
|
||||
ckpt = tf.train.Checkpoint(generator=generator)
|
||||
out_dir = os.path.join(ckpt_dir, 'g_clone' if is_g_clone else 'generator')
|
||||
manager = tf.train.CheckpointManager(ckpt, out_dir, max_to_keep=1)
|
||||
manager.save(checkpoint_number=0)
|
||||
return
|
||||
|
||||
|
||||
def convert_official_discriminator_weights(ckpt_dir, use_custom_cuda):
|
||||
discriminator = load_discriminator(d_params=None, ckpt_dir=None, custom_cuda=use_custom_cuda)
|
||||
|
||||
# restore official ones
|
||||
official_checkpoint = tf.train.latest_checkpoint('./official-pretrained')
|
||||
official_vars = tf.train.list_variables(official_checkpoint)
|
||||
|
||||
# get name mapper
|
||||
name_mapper = variable_name_mapper_d(discriminator)
|
||||
for name_d, tvar in name_mapper.items():
|
||||
print(f'{name_d}: {tvar.name}')
|
||||
|
||||
# check shape
|
||||
check_shape(name_mapper, official_vars)
|
||||
|
||||
# restore
|
||||
tf.compat.v1.train.init_from_checkpoint(official_checkpoint, assignment_map=name_mapper)
|
||||
|
||||
# save
|
||||
ckpt = tf.train.Checkpoint(discriminator=discriminator)
|
||||
out_dir = os.path.join(ckpt_dir, 'discriminator')
|
||||
manager = tf.train.CheckpointManager(ckpt, out_dir, max_to_keep=1)
|
||||
manager.save(checkpoint_number=0)
|
||||
return
|
||||
|
||||
|
||||
def convert_official_weights_together(ckpt_dir, use_custom_cuda):
|
||||
# instantiate all models
|
||||
discriminator = load_discriminator(d_params=None, ckpt_dir=None, custom_cuda=use_custom_cuda)
|
||||
generator = load_generator(g_params=None, is_g_clone=False, ckpt_dir=None, custom_cuda=use_custom_cuda)
|
||||
g_clone = load_generator(g_params=None, is_g_clone=True, ckpt_dir=None, custom_cuda=use_custom_cuda)
|
||||
|
||||
# restore official ones
|
||||
official_checkpoint = tf.train.latest_checkpoint('./official-pretrained')
|
||||
official_vars = tf.train.list_variables(official_checkpoint)
|
||||
for name, shape in official_vars:
|
||||
print(f'{name}: {shape}')
|
||||
|
||||
# get name mapper
|
||||
name_mapper_d = variable_name_mapper_d(discriminator)
|
||||
name_mapper_g1 = variable_name_mapper_g(generator, is_g_clone=False)
|
||||
name_mapper_g2 = variable_name_mapper_g(g_clone, is_g_clone=True)
|
||||
name_mapper = {**name_mapper_d, **name_mapper_g1, **name_mapper_g2}
|
||||
|
||||
# check shape
|
||||
check_shape(name_mapper, official_vars)
|
||||
|
||||
# restore
|
||||
tf.compat.v1.train.init_from_checkpoint(official_checkpoint, assignment_map=name_mapper)
|
||||
|
||||
# save
|
||||
ckpt = tf.train.Checkpoint(discriminator=discriminator,
|
||||
generator=generator,
|
||||
g_clone=g_clone)
|
||||
manager = tf.train.CheckpointManager(ckpt, ckpt_dir, max_to_keep=1)
|
||||
manager.save(checkpoint_number=0)
|
||||
return
|
||||
|
||||
|
||||
def main():
|
||||
from tf_utils import allow_memory_growth
|
||||
|
||||
allow_memory_growth()
|
||||
|
||||
ckpt_dir_base = './official-converted'
|
||||
for use_custom_cuda in [True, False]:
|
||||
ckpt_dir = os.path.join(ckpt_dir_base, 'cuda') if use_custom_cuda else os.path.join(ckpt_dir_base, 'ref')
|
||||
|
||||
convert_official_weights_together(ckpt_dir, use_custom_cuda)
|
||||
|
||||
# convert_official_discriminator_weights(ckpt_dir, use_custom_cuda)
|
||||
# for is_g_clone in [True, False]:
|
||||
# convert_official_generator_weights(ckpt_dir, is_g_clone, use_custom_cuda)
|
||||
return
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user