-
Notifications
You must be signed in to change notification settings - Fork 19
/
legacy.py
executable file
·117 lines (96 loc) · 5.13 KB
/
legacy.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
#
# This work is licensed under the Creative Commons Attribution-NonCommercial
# 4.0 International License. To view a copy of this license, visit
# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
import pickle
import inspect
import numpy as np
import tfutil
import networks
#----------------------------------------------------------------------------
# Custom unpickler that is able to load network pickles produced by
# the old Theano implementation.
class LegacyUnpickler(pickle.Unpickler):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def find_class(self, module, name):
if module == 'network' and name == 'Network':
return tfutil.Network
return super().find_class(module, name)
#----------------------------------------------------------------------------
# Import handler for tfutil.Network that silently converts networks produced
# by the old Theano implementation to a suitable format.
theano_gan_remap = {
'G_paper': 'G_paper',
'G_progressive_8': 'G_paper',
'D_paper': 'D_paper',
'D_progressive_8': 'D_paper'}
def patch_theano_gan(state):
if 'version' in state or state['build_func_spec']['func'] not in theano_gan_remap:
return state
spec = dict(state['build_func_spec'])
func = spec.pop('func')
resolution = spec.get('resolution', 32)
resolution_log2 = int(np.log2(resolution))
use_wscale = spec.get('use_wscale', True)
assert spec.pop('label_size', 0) == 0
assert spec.pop('use_batchnorm', False) == False
assert spec.pop('tanh_at_end', None) is None
assert spec.pop('mbstat_func', 'Tstdeps') == 'Tstdeps'
assert spec.pop('mbstat_avg', 'all') == 'all'
assert spec.pop('mbdisc_kernels', None) is None
spec.pop( 'use_gdrop', True) # doesn't make a difference
assert spec.pop('use_layernorm', False) == False
spec[ 'fused_scale'] = False
spec[ 'mbstd_group_size'] = 16
vars = []
param_iter = iter(state['param_values'])
relu = np.sqrt(2); linear = 1.0
def flatten2(w): return w.reshape(w.shape[0], -1)
def he_std(gain, w): return gain / np.sqrt(np.prod(w.shape[:-1]))
def wscale(gain, w): return w * next(param_iter) / he_std(gain, w) if use_wscale else w
def layer(name, gain, w): return [(name + '/weight', wscale(gain, w)), (name + '/bias', next(param_iter))]
if func.startswith('G'):
vars += layer('4x4/Dense', relu/4, flatten2(next(param_iter).transpose(1,0,2,3)))
vars += layer('4x4/Conv', relu, next(param_iter).transpose(2,3,1,0)[::-1,::-1])
for res in range(3, resolution_log2 + 1):
vars += layer('%dx%d/Conv0' % (2**res, 2**res), relu, next(param_iter).transpose(2,3,1,0)[::-1,::-1])
vars += layer('%dx%d/Conv1' % (2**res, 2**res), relu, next(param_iter).transpose(2,3,1,0)[::-1,::-1])
for lod in range(0, resolution_log2 - 1):
vars += layer('ToRGB_lod%d' % lod, linear, next(param_iter)[np.newaxis, np.newaxis])
if func.startswith('D'):
vars += layer('FromRGB_lod0', relu, next(param_iter)[np.newaxis, np.newaxis])
for res in range(resolution_log2, 2, -1):
vars += layer('%dx%d/Conv0' % (2**res, 2**res), relu, next(param_iter).transpose(2,3,1,0)[::-1,::-1])
vars += layer('%dx%d/Conv1' % (2**res, 2**res), relu, next(param_iter).transpose(2,3,1,0)[::-1,::-1])
vars += layer('FromRGB_lod%d' % (resolution_log2 - (res - 1)), relu, next(param_iter)[np.newaxis, np.newaxis])
vars += layer('4x4/Conv', relu, next(param_iter).transpose(2,3,1,0)[::-1,::-1])
vars += layer('4x4/Dense0', relu, flatten2(next(param_iter)[:,:,::-1,::-1]).transpose())
vars += layer('4x4/Dense1', linear, next(param_iter))
vars += [('lod', state['toplevel_params']['cur_lod'])]
return {
'version': 2,
'name': func,
'build_module_src': inspect.getsource(networks),
'build_func_name': theano_gan_remap[func],
'static_kwargs': spec,
'variables': vars}
tfutil.network_import_handlers.append(patch_theano_gan)
#----------------------------------------------------------------------------
# Import handler for tfutil.Network that ignores unsupported/deprecated
# networks produced by older versions of the code.
def ignore_unknown_theano_network(state):
if 'version' in state:
return state
print('Ignoring unknown Theano network:', state['build_func_spec']['func'])
return {
'version': 2,
'name': 'Dummy',
'build_module_src': 'def dummy(input, **kwargs): input.set_shape([None, 1]); return input',
'build_func_name': 'dummy',
'static_kwargs': {},
'variables': []}
tfutil.network_import_handlers.append(ignore_unknown_theano_network)
#----------------------------------------------------------------------------