Diggz10 commited on
Commit
6e9b99b
·
verified ·
1 Parent(s): 5a81e39

Delete training

Browse files
training/__init__.py DELETED
@@ -1,7 +0,0 @@
1
- # Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
2
- #
3
- # This work is made available under the Nvidia Source Code License-NC.
4
- # To view a copy of this license, visit
5
- # https://nvlabs.github.io/stylegan2/license.html
6
-
7
- # empty
 
 
 
 
 
 
 
 
training/dataset.py DELETED
@@ -1,199 +0,0 @@
1
- # Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
2
- #
3
- # This work is made available under the Nvidia Source Code License-NC.
4
- # To view a copy of this license, visit
5
- # https://nvlabs.github.io/stylegan2/license.html
6
-
7
- """Multi-resolution input data pipeline."""
8
-
9
- import os
10
- import glob
11
- import numpy as np
12
- import tensorflow as tf
13
- import dnnlib
14
- import dnnlib.tflib as tflib
15
-
16
- #----------------------------------------------------------------------------
17
- # Dataset class that loads data from tfrecords files.
18
-
19
- class TFRecordDataset:
20
- def __init__(self,
21
- tfrecord_dir, # Directory containing a collection of tfrecords files.
22
- resolution = None, # Dataset resolution, None = autodetect.
23
- label_file = None, # Relative path of the labels file, None = autodetect.
24
- max_label_size = 0, # 0 = no labels, 'full' = full labels, <int> = N first label components.
25
- max_images = None, # Maximum number of images to use, None = use all images.
26
- repeat = True, # Repeat dataset indefinitely?
27
- shuffle_mb = 4096, # Shuffle data within specified window (megabytes), 0 = disable shuffling.
28
- prefetch_mb = 2048, # Amount of data to prefetch (megabytes), 0 = disable prefetching.
29
- buffer_mb = 256, # Read buffer size (megabytes).
30
- num_threads = 2): # Number of concurrent threads.
31
-
32
- self.tfrecord_dir = tfrecord_dir
33
- self.resolution = None
34
- self.resolution_log2 = None
35
- self.shape = [] # [channels, height, width]
36
- self.dtype = 'uint8'
37
- self.dynamic_range = [0, 255]
38
- self.label_file = label_file
39
- self.label_size = None # components
40
- self.label_dtype = None
41
- self._np_labels = None
42
- self._tf_minibatch_in = None
43
- self._tf_labels_var = None
44
- self._tf_labels_dataset = None
45
- self._tf_datasets = dict()
46
- self._tf_iterator = None
47
- self._tf_init_ops = dict()
48
- self._tf_minibatch_np = None
49
- self._cur_minibatch = -1
50
- self._cur_lod = -1
51
-
52
- # List tfrecords files and inspect their shapes.
53
- assert os.path.isdir(self.tfrecord_dir)
54
- tfr_files = sorted(glob.glob(os.path.join(self.tfrecord_dir, '*.tfrecords')))
55
- assert len(tfr_files) >= 1
56
- tfr_shapes = []
57
- for tfr_file in tfr_files:
58
- tfr_opt = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.NONE)
59
- for record in tf.python_io.tf_record_iterator(tfr_file, tfr_opt):
60
- tfr_shapes.append(self.parse_tfrecord_np(record).shape)
61
- break
62
-
63
- # Autodetect label filename.
64
- if self.label_file is None:
65
- guess = sorted(glob.glob(os.path.join(self.tfrecord_dir, '*.labels')))
66
- if len(guess):
67
- self.label_file = guess[0]
68
- elif not os.path.isfile(self.label_file):
69
- guess = os.path.join(self.tfrecord_dir, self.label_file)
70
- if os.path.isfile(guess):
71
- self.label_file = guess
72
-
73
- # Determine shape and resolution.
74
- max_shape = max(tfr_shapes, key=np.prod)
75
- self.resolution = resolution if resolution is not None else max_shape[1]
76
- self.resolution_log2 = int(np.log2(self.resolution))
77
- self.shape = [max_shape[0], self.resolution, self.resolution]
78
- tfr_lods = [self.resolution_log2 - int(np.log2(shape[1])) for shape in tfr_shapes]
79
- assert all(shape[0] == max_shape[0] for shape in tfr_shapes)
80
- assert all(shape[1] == shape[2] for shape in tfr_shapes)
81
- assert all(shape[1] == self.resolution // (2**lod) for shape, lod in zip(tfr_shapes, tfr_lods))
82
- assert all(lod in tfr_lods for lod in range(self.resolution_log2 - 1))
83
-
84
- # Load labels.
85
- assert max_label_size == 'full' or max_label_size >= 0
86
- self._np_labels = np.zeros([1<<30, 0], dtype=np.float32)
87
- if self.label_file is not None and max_label_size != 0:
88
- self._np_labels = np.load(self.label_file)
89
- assert self._np_labels.ndim == 2
90
- if max_label_size != 'full' and self._np_labels.shape[1] > max_label_size:
91
- self._np_labels = self._np_labels[:, :max_label_size]
92
- if max_images is not None and self._np_labels.shape[0] > max_images:
93
- self._np_labels = self._np_labels[:max_images]
94
- self.label_size = self._np_labels.shape[1]
95
- self.label_dtype = self._np_labels.dtype.name
96
-
97
- # Build TF expressions.
98
- with tf.name_scope('Dataset'), tf.device('/cpu:0'):
99
- self._tf_minibatch_in = tf.placeholder(tf.int64, name='minibatch_in', shape=[])
100
- self._tf_labels_var = tflib.create_var_with_large_initial_value(self._np_labels, name='labels_var')
101
- self._tf_labels_dataset = tf.data.Dataset.from_tensor_slices(self._tf_labels_var)
102
- for tfr_file, tfr_shape, tfr_lod in zip(tfr_files, tfr_shapes, tfr_lods):
103
- if tfr_lod < 0:
104
- continue
105
- dset = tf.data.TFRecordDataset(tfr_file, compression_type='', buffer_size=buffer_mb<<20)
106
- if max_images is not None:
107
- dset = dset.take(max_images)
108
- dset = dset.map(self.parse_tfrecord_tf, num_parallel_calls=num_threads)
109
- dset = tf.data.Dataset.zip((dset, self._tf_labels_dataset))
110
- bytes_per_item = np.prod(tfr_shape) * np.dtype(self.dtype).itemsize
111
- if shuffle_mb > 0:
112
- dset = dset.shuffle(((shuffle_mb << 20) - 1) // bytes_per_item + 1)
113
- if repeat:
114
- dset = dset.repeat()
115
- if prefetch_mb > 0:
116
- dset = dset.prefetch(((prefetch_mb << 20) - 1) // bytes_per_item + 1)
117
- dset = dset.batch(self._tf_minibatch_in)
118
- self._tf_datasets[tfr_lod] = dset
119
- self._tf_iterator = tf.data.Iterator.from_structure(self._tf_datasets[0].output_types, self._tf_datasets[0].output_shapes)
120
- self._tf_init_ops = {lod: self._tf_iterator.make_initializer(dset) for lod, dset in self._tf_datasets.items()}
121
-
122
- def close(self):
123
- pass
124
-
125
- # Use the given minibatch size and level-of-detail for the data returned by get_minibatch_tf().
126
- def configure(self, minibatch_size, lod=0):
127
- lod = int(np.floor(lod))
128
- assert minibatch_size >= 1 and lod in self._tf_datasets
129
- if self._cur_minibatch != minibatch_size or self._cur_lod != lod:
130
- self._tf_init_ops[lod].run({self._tf_minibatch_in: minibatch_size})
131
- self._cur_minibatch = minibatch_size
132
- self._cur_lod = lod
133
-
134
- # Get next minibatch as TensorFlow expressions.
135
- def get_minibatch_tf(self): # => images, labels
136
- return self._tf_iterator.get_next()
137
-
138
- # Get next minibatch as NumPy arrays.
139
- def get_minibatch_np(self, minibatch_size, lod=0): # => images, labels
140
- self.configure(minibatch_size, lod)
141
- with tf.name_scope('Dataset'):
142
- if self._tf_minibatch_np is None:
143
- self._tf_minibatch_np = self.get_minibatch_tf()
144
- return tflib.run(self._tf_minibatch_np)
145
-
146
- # Get random labels as TensorFlow expression.
147
- def get_random_labels_tf(self, minibatch_size): # => labels
148
- with tf.name_scope('Dataset'):
149
- if self.label_size > 0:
150
- with tf.device('/cpu:0'):
151
- return tf.gather(self._tf_labels_var, tf.random_uniform([minibatch_size], 0, self._np_labels.shape[0], dtype=tf.int32))
152
- return tf.zeros([minibatch_size, 0], self.label_dtype)
153
-
154
- # Get random labels as NumPy array.
155
- def get_random_labels_np(self, minibatch_size): # => labels
156
- if self.label_size > 0:
157
- return self._np_labels[np.random.randint(self._np_labels.shape[0], size=[minibatch_size])]
158
- return np.zeros([minibatch_size, 0], self.label_dtype)
159
-
160
- # Parse individual image from a tfrecords file into TensorFlow expression.
161
- @staticmethod
162
- def parse_tfrecord_tf(record):
163
- features = tf.parse_single_example(record, features={
164
- 'shape': tf.FixedLenFeature([3], tf.int64),
165
- 'data': tf.FixedLenFeature([], tf.string)})
166
- data = tf.decode_raw(features['data'], tf.uint8)
167
- return tf.reshape(data, features['shape'])
168
-
169
- # Parse individual image from a tfrecords file into NumPy array.
170
- @staticmethod
171
- def parse_tfrecord_np(record):
172
- ex = tf.train.Example()
173
- ex.ParseFromString(record)
174
- shape = ex.features.feature['shape'].int64_list.value # pylint: disable=no-member
175
- data = ex.features.feature['data'].bytes_list.value[0] # pylint: disable=no-member
176
- return np.fromstring(data, np.uint8).reshape(shape)
177
-
178
- #----------------------------------------------------------------------------
179
- # Helper func for constructing a dataset object using the given options.
180
-
181
- def load_dataset(class_name=None, data_dir=None, verbose=False, **kwargs):
182
- kwargs = dict(kwargs)
183
- if 'tfrecord_dir' in kwargs:
184
- if class_name is None:
185
- class_name = __name__ + '.TFRecordDataset'
186
- if data_dir is not None:
187
- kwargs['tfrecord_dir'] = os.path.join(data_dir, kwargs['tfrecord_dir'])
188
-
189
- assert class_name is not None
190
- if verbose:
191
- print('Streaming data using %s...' % class_name)
192
- dataset = dnnlib.util.get_obj_by_name(class_name)(**kwargs)
193
- if verbose:
194
- print('Dataset shape =', np.int32(dataset.shape).tolist())
195
- print('Dynamic range =', dataset.dynamic_range)
196
- print('Label size =', dataset.label_size)
197
- return dataset
198
-
199
- #----------------------------------------------------------------------------
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
training/loss.py DELETED
@@ -1,197 +0,0 @@
1
- # Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
2
- #
3
- # This work is made available under the Nvidia Source Code License-NC.
4
- # To view a copy of this license, visit
5
- # https://nvlabs.github.io/stylegan2/license.html
6
-
7
- """Loss functions."""
8
-
9
- import numpy as np
10
- import tensorflow as tf
11
- import dnnlib.tflib as tflib
12
- from dnnlib.tflib.autosummary import autosummary
13
-
14
- #----------------------------------------------------------------------------
15
- # Logistic loss from the paper
16
- # "Generative Adversarial Nets", Goodfellow et al. 2014
17
-
18
- def G_logistic(G, D, opt, training_set, minibatch_size):
19
- _ = opt
20
- latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:])
21
- labels = training_set.get_random_labels_tf(minibatch_size)
22
- fake_images_out = G.get_output_for(latents, labels, is_training=True)
23
- fake_scores_out = D.get_output_for(fake_images_out, labels, is_training=True)
24
- loss = -tf.nn.softplus(fake_scores_out) # log(1-sigmoid(fake_scores_out)) # pylint: disable=invalid-unary-operand-type
25
- return loss, None
26
-
27
- def G_logistic_ns(G, D, opt, training_set, minibatch_size):
28
- _ = opt
29
- latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:])
30
- labels = training_set.get_random_labels_tf(minibatch_size)
31
- fake_images_out = G.get_output_for(latents, labels, is_training=True)
32
- fake_scores_out = D.get_output_for(fake_images_out, labels, is_training=True)
33
- loss = tf.nn.softplus(-fake_scores_out) # -log(sigmoid(fake_scores_out))
34
- return loss, None
35
-
36
- def D_logistic(G, D, opt, training_set, minibatch_size, reals, labels):
37
- _ = opt, training_set
38
- latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:])
39
- fake_images_out = G.get_output_for(latents, labels, is_training=True)
40
- real_scores_out = D.get_output_for(reals, labels, is_training=True)
41
- fake_scores_out = D.get_output_for(fake_images_out, labels, is_training=True)
42
- real_scores_out = autosummary('Loss/scores/real', real_scores_out)
43
- fake_scores_out = autosummary('Loss/scores/fake', fake_scores_out)
44
- loss = tf.nn.softplus(fake_scores_out) # -log(1-sigmoid(fake_scores_out))
45
- loss += tf.nn.softplus(-real_scores_out) # -log(sigmoid(real_scores_out)) # pylint: disable=invalid-unary-operand-type
46
- return loss, None
47
-
48
- #----------------------------------------------------------------------------
49
- # R1 and R2 regularizers from the paper
50
- # "Which Training Methods for GANs do actually Converge?", Mescheder et al. 2018
51
-
52
- def D_logistic_r1(G, D, opt, training_set, minibatch_size, reals, labels, gamma=10.0):
53
- _ = opt, training_set
54
- latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:])
55
- fake_images_out = G.get_output_for(latents, labels, is_training=True)
56
- real_scores_out = D.get_output_for(reals, labels, is_training=True)
57
- fake_scores_out = D.get_output_for(fake_images_out, labels, is_training=True)
58
- real_scores_out = autosummary('Loss/scores/real', real_scores_out)
59
- fake_scores_out = autosummary('Loss/scores/fake', fake_scores_out)
60
- loss = tf.nn.softplus(fake_scores_out) # -log(1-sigmoid(fake_scores_out))
61
- loss += tf.nn.softplus(-real_scores_out) # -log(sigmoid(real_scores_out)) # pylint: disable=invalid-unary-operand-type
62
-
63
- with tf.name_scope('GradientPenalty'):
64
- real_grads = tf.gradients(tf.reduce_sum(real_scores_out), [reals])[0]
65
- gradient_penalty = tf.reduce_sum(tf.square(real_grads), axis=[1,2,3])
66
- gradient_penalty = autosummary('Loss/gradient_penalty', gradient_penalty)
67
- reg = gradient_penalty * (gamma * 0.5)
68
- return loss, reg
69
-
70
- def D_logistic_r2(G, D, opt, training_set, minibatch_size, reals, labels, gamma=10.0):
71
- _ = opt, training_set
72
- latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:])
73
- fake_images_out = G.get_output_for(latents, labels, is_training=True)
74
- real_scores_out = D.get_output_for(reals, labels, is_training=True)
75
- fake_scores_out = D.get_output_for(fake_images_out, labels, is_training=True)
76
- real_scores_out = autosummary('Loss/scores/real', real_scores_out)
77
- fake_scores_out = autosummary('Loss/scores/fake', fake_scores_out)
78
- loss = tf.nn.softplus(fake_scores_out) # -log(1-sigmoid(fake_scores_out))
79
- loss += tf.nn.softplus(-real_scores_out) # -log(sigmoid(real_scores_out)) # pylint: disable=invalid-unary-operand-type
80
-
81
- with tf.name_scope('GradientPenalty'):
82
- fake_grads = tf.gradients(tf.reduce_sum(fake_scores_out), [fake_images_out])[0]
83
- gradient_penalty = tf.reduce_sum(tf.square(fake_grads), axis=[1,2,3])
84
- gradient_penalty = autosummary('Loss/gradient_penalty', gradient_penalty)
85
- reg = gradient_penalty * (gamma * 0.5)
86
- return loss, reg
87
-
88
- #----------------------------------------------------------------------------
89
- # WGAN loss from the paper
90
- # "Wasserstein Generative Adversarial Networks", Arjovsky et al. 2017
91
-
92
- def G_wgan(G, D, opt, training_set, minibatch_size):
93
- _ = opt
94
- latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:])
95
- labels = training_set.get_random_labels_tf(minibatch_size)
96
- fake_images_out = G.get_output_for(latents, labels, is_training=True)
97
- fake_scores_out = D.get_output_for(fake_images_out, labels, is_training=True)
98
- loss = -fake_scores_out
99
- return loss, None
100
-
101
- def D_wgan(G, D, opt, training_set, minibatch_size, reals, labels, wgan_epsilon=0.001):
102
- _ = opt, training_set
103
- latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:])
104
- fake_images_out = G.get_output_for(latents, labels, is_training=True)
105
- real_scores_out = D.get_output_for(reals, labels, is_training=True)
106
- fake_scores_out = D.get_output_for(fake_images_out, labels, is_training=True)
107
- real_scores_out = autosummary('Loss/scores/real', real_scores_out)
108
- fake_scores_out = autosummary('Loss/scores/fake', fake_scores_out)
109
- loss = fake_scores_out - real_scores_out
110
- with tf.name_scope('EpsilonPenalty'):
111
- epsilon_penalty = autosummary('Loss/epsilon_penalty', tf.square(real_scores_out))
112
- loss += epsilon_penalty * wgan_epsilon
113
- return loss, None
114
-
115
- #----------------------------------------------------------------------------
116
- # WGAN-GP loss from the paper
117
- # "Improved Training of Wasserstein GANs", Gulrajani et al. 2017
118
-
119
- def D_wgan_gp(G, D, opt, training_set, minibatch_size, reals, labels, wgan_lambda=10.0, wgan_epsilon=0.001, wgan_target=1.0):
120
- _ = opt, training_set
121
- latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:])
122
- fake_images_out = G.get_output_for(latents, labels, is_training=True)
123
- real_scores_out = D.get_output_for(reals, labels, is_training=True)
124
- fake_scores_out = D.get_output_for(fake_images_out, labels, is_training=True)
125
- real_scores_out = autosummary('Loss/scores/real', real_scores_out)
126
- fake_scores_out = autosummary('Loss/scores/fake', fake_scores_out)
127
- loss = fake_scores_out - real_scores_out
128
- with tf.name_scope('EpsilonPenalty'):
129
- epsilon_penalty = autosummary('Loss/epsilon_penalty', tf.square(real_scores_out))
130
- loss += epsilon_penalty * wgan_epsilon
131
-
132
- with tf.name_scope('GradientPenalty'):
133
- mixing_factors = tf.random_uniform([minibatch_size, 1, 1, 1], 0.0, 1.0, dtype=fake_images_out.dtype)
134
- mixed_images_out = tflib.lerp(tf.cast(reals, fake_images_out.dtype), fake_images_out, mixing_factors)
135
- mixed_scores_out = D.get_output_for(mixed_images_out, labels, is_training=True)
136
- mixed_scores_out = autosummary('Loss/scores/mixed', mixed_scores_out)
137
- mixed_grads = tf.gradients(tf.reduce_sum(mixed_scores_out), [mixed_images_out])[0]
138
- mixed_norms = tf.sqrt(tf.reduce_sum(tf.square(mixed_grads), axis=[1,2,3]))
139
- mixed_norms = autosummary('Loss/mixed_norms', mixed_norms)
140
- gradient_penalty = tf.square(mixed_norms - wgan_target)
141
- reg = gradient_penalty * (wgan_lambda / (wgan_target**2))
142
- return loss, reg
143
-
144
- #----------------------------------------------------------------------------
145
- # Non-saturating logistic loss with path length regularizer from the paper
146
- # "Analyzing and Improving the Image Quality of StyleGAN", Karras et al. 2019
147
-
148
- def G_logistic_ns_pathreg(G, D, opt, training_set, minibatch_size, pl_minibatch_shrink=2, pl_decay=0.01, pl_weight=2.0):
149
- _ = opt
150
- latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:])
151
- labels = training_set.get_random_labels_tf(minibatch_size)
152
- fake_images_out, fake_dlatents_out = G.get_output_for(latents, labels, is_training=True, return_dlatents=True)
153
- fake_scores_out = D.get_output_for(fake_images_out, labels, is_training=True)
154
- loss = tf.nn.softplus(-fake_scores_out) # -log(sigmoid(fake_scores_out))
155
-
156
- # Path length regularization.
157
- with tf.name_scope('PathReg'):
158
-
159
- # Evaluate the regularization term using a smaller minibatch to conserve memory.
160
- if pl_minibatch_shrink > 1:
161
- pl_minibatch = minibatch_size // pl_minibatch_shrink
162
- pl_latents = tf.random_normal([pl_minibatch] + G.input_shapes[0][1:])
163
- pl_labels = training_set.get_random_labels_tf(pl_minibatch)
164
- fake_images_out, fake_dlatents_out = G.get_output_for(pl_latents, pl_labels, is_training=True, return_dlatents=True)
165
-
166
- # Compute |J*y|.
167
- pl_noise = tf.random_normal(tf.shape(fake_images_out)) / np.sqrt(np.prod(G.output_shape[2:]))
168
- pl_grads = tf.gradients(tf.reduce_sum(fake_images_out * pl_noise), [fake_dlatents_out])[0]
169
- pl_lengths = tf.sqrt(tf.reduce_mean(tf.reduce_sum(tf.square(pl_grads), axis=2), axis=1))
170
- pl_lengths = autosummary('Loss/pl_lengths', pl_lengths)
171
-
172
- # Track exponential moving average of |J*y|.
173
- with tf.control_dependencies(None):
174
- pl_mean_var = tf.Variable(name='pl_mean', trainable=False, initial_value=0.0, dtype=tf.float32)
175
- pl_mean = pl_mean_var + pl_decay * (tf.reduce_mean(pl_lengths) - pl_mean_var)
176
- pl_update = tf.assign(pl_mean_var, pl_mean)
177
-
178
- # Calculate (|J*y|-a)^2.
179
- with tf.control_dependencies([pl_update]):
180
- pl_penalty = tf.square(pl_lengths - pl_mean)
181
- pl_penalty = autosummary('Loss/pl_penalty', pl_penalty)
182
-
183
- # Apply weight.
184
- #
185
- # Note: The division in pl_noise decreases the weight by num_pixels, and the reduce_mean
186
- # in pl_lengths decreases it by num_affine_layers. The effective weight then becomes:
187
- #
188
- # gamma_pl = pl_weight / num_pixels / num_affine_layers
189
- # = 2 / (r^2) / (log2(r) * 2 - 2)
190
- # = 1 / (r^2 * (log2(r) - 1))
191
- # = ln(2) / (r^2 * (ln(r) - ln(2))
192
- #
193
- reg = pl_penalty * pl_weight
194
-
195
- return loss, reg
196
-
197
- #----------------------------------------------------------------------------
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
training/misc.py DELETED
@@ -1,145 +0,0 @@
1
- # Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
2
- #
3
- # This work is made available under the Nvidia Source Code License-NC.
4
- # To view a copy of this license, visit
5
- # https://nvlabs.github.io/stylegan2/license.html
6
-
7
- """Miscellaneous utility functions."""
8
-
9
- import os
10
- import pickle
11
- import numpy as np
12
- import PIL.Image
13
- import PIL.ImageFont
14
- import dnnlib
15
-
16
- #----------------------------------------------------------------------------
17
- # Convenience wrappers for pickle that are able to load data produced by
18
- # older versions of the code, and from external URLs.
19
-
20
- def open_file_or_url(file_or_url):
21
- if dnnlib.util.is_url(file_or_url):
22
- return dnnlib.util.open_url(file_or_url, cache_dir='models')
23
- return open(file_or_url, 'rb')
24
-
25
- def load_pkl(file_or_url):
26
- with open_file_or_url(file_or_url) as file:
27
- return pickle.load(file, encoding='latin1')
28
-
29
- def save_pkl(obj, filename):
30
- with open(filename, 'wb') as file:
31
- pickle.dump(obj, file, protocol=pickle.HIGHEST_PROTOCOL)
32
-
33
- #----------------------------------------------------------------------------
34
- # Image utils.
35
-
36
- def adjust_dynamic_range(data, drange_in, drange_out):
37
- if drange_in != drange_out:
38
- scale = (np.float32(drange_out[1]) - np.float32(drange_out[0])) / (np.float32(drange_in[1]) - np.float32(drange_in[0]))
39
- bias = (np.float32(drange_out[0]) - np.float32(drange_in[0]) * scale)
40
- data = data * scale + bias
41
- return data
42
-
43
- def create_image_grid(images, grid_size=None):
44
- assert images.ndim == 3 or images.ndim == 4
45
- num, img_w, img_h = images.shape[0], images.shape[-1], images.shape[-2]
46
-
47
- if grid_size is not None:
48
- grid_w, grid_h = tuple(grid_size)
49
- else:
50
- grid_w = max(int(np.ceil(np.sqrt(num))), 1)
51
- grid_h = max((num - 1) // grid_w + 1, 1)
52
-
53
- grid = np.zeros(list(images.shape[1:-2]) + [grid_h * img_h, grid_w * img_w], dtype=images.dtype)
54
- for idx in range(num):
55
- x = (idx % grid_w) * img_w
56
- y = (idx // grid_w) * img_h
57
- grid[..., y : y + img_h, x : x + img_w] = images[idx]
58
- return grid
59
-
60
- def convert_to_pil_image(image, drange=[0,1]):
61
- assert image.ndim == 2 or image.ndim == 3
62
- if image.ndim == 3:
63
- if image.shape[0] == 1:
64
- image = image[0] # grayscale CHW => HW
65
- else:
66
- image = image.transpose(1, 2, 0) # CHW -> HWC
67
-
68
- image = adjust_dynamic_range(image, drange, [0,255])
69
- image = np.rint(image).clip(0, 255).astype(np.uint8)
70
- fmt = 'RGB' if image.ndim == 3 else 'L'
71
- return PIL.Image.fromarray(image, fmt)
72
-
73
- def save_image_grid(images, filename, drange=[0,1], grid_size=None):
74
- convert_to_pil_image(create_image_grid(images, grid_size), drange).save(filename)
75
-
76
- def apply_mirror_augment(minibatch):
77
- mask = np.random.rand(minibatch.shape[0]) < 0.5
78
- minibatch = np.array(minibatch)
79
- minibatch[mask] = minibatch[mask, :, :, ::-1]
80
- return minibatch
81
-
82
- #----------------------------------------------------------------------------
83
- # Loading data from previous training runs.
84
-
85
- def parse_config_for_previous_run(run_dir):
86
- with open(os.path.join(run_dir, 'submit_config.pkl'), 'rb') as f:
87
- data = pickle.load(f)
88
- data = data.get('run_func_kwargs', {})
89
- return dict(train=data, dataset=data.get('dataset_args', {}))
90
-
91
- #----------------------------------------------------------------------------
92
- # Size and contents of the image snapshot grids that are exported
93
- # periodically during training.
94
-
95
- def setup_snapshot_image_grid(training_set,
96
- size = '1080p', # '1080p' = to be viewed on 1080p display, '4k' = to be viewed on 4k display.
97
- layout = 'random'): # 'random' = grid contents are selected randomly, 'row_per_class' = each row corresponds to one class label.
98
-
99
- # Select size.
100
- gw = 1; gh = 1
101
- if size == '1080p':
102
- gw = np.clip(1920 // training_set.shape[2], 3, 32)
103
- gh = np.clip(1080 // training_set.shape[1], 2, 32)
104
- if size == '4k':
105
- gw = np.clip(3840 // training_set.shape[2], 7, 32)
106
- gh = np.clip(2160 // training_set.shape[1], 4, 32)
107
- if size == '8k':
108
- gw = np.clip(7680 // training_set.shape[2], 7, 32)
109
- gh = np.clip(4320 // training_set.shape[1], 4, 32)
110
-
111
- # Initialize data arrays.
112
- reals = np.zeros([gw * gh] + training_set.shape, dtype=training_set.dtype)
113
- labels = np.zeros([gw * gh, training_set.label_size], dtype=training_set.label_dtype)
114
-
115
- # Random layout.
116
- if layout == 'random':
117
- reals[:], labels[:] = training_set.get_minibatch_np(gw * gh)
118
-
119
- # Class-conditional layouts.
120
- class_layouts = dict(row_per_class=[gw,1], col_per_class=[1,gh], class4x4=[4,4])
121
- if layout in class_layouts:
122
- bw, bh = class_layouts[layout]
123
- nw = (gw - 1) // bw + 1
124
- nh = (gh - 1) // bh + 1
125
- blocks = [[] for _i in range(nw * nh)]
126
- for _iter in range(1000000):
127
- real, label = training_set.get_minibatch_np(1)
128
- idx = np.argmax(label[0])
129
- while idx < len(blocks) and len(blocks[idx]) >= bw * bh:
130
- idx += training_set.label_size
131
- if idx < len(blocks):
132
- blocks[idx].append((real, label))
133
- if all(len(block) >= bw * bh for block in blocks):
134
- break
135
- for i, block in enumerate(blocks):
136
- for j, (real, label) in enumerate(block):
137
- x = (i % nw) * bw + j % bw
138
- y = (i // nw) * bh + j // bw
139
- if x < gw and y < gh:
140
- reals[x + y * gw] = real[0]
141
- labels[x + y * gw] = label[0]
142
-
143
- return (gw, gh), reals, labels
144
-
145
- #----------------------------------------------------------------------------
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
training/networks_stylegan.py DELETED
@@ -1,660 +0,0 @@
1
- # Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
2
- #
3
- # This work is made available under the Nvidia Source Code License-NC.
4
- # To view a copy of this license, visit
5
- # https://nvlabs.github.io/stylegan2/license.html
6
-
7
- """Network architectures used in the StyleGAN paper."""
8
-
9
- import numpy as np
10
- import tensorflow as tf
11
- import dnnlib
12
- import dnnlib.tflib as tflib
13
-
14
- # NOTE: Do not import any application-specific modules here!
15
- # Specify all network parameters as kwargs.
16
-
17
- #----------------------------------------------------------------------------
18
- # Primitive ops for manipulating 4D activation tensors.
19
- # The gradients of these are not necessary efficient or even meaningful.
20
-
21
- def _blur2d(x, f=[1,2,1], normalize=True, flip=False, stride=1):
22
- assert x.shape.ndims == 4 and all(dim.value is not None for dim in x.shape[1:])
23
- assert isinstance(stride, int) and stride >= 1
24
-
25
- # Finalize filter kernel.
26
- f = np.array(f, dtype=np.float32)
27
- if f.ndim == 1:
28
- f = f[:, np.newaxis] * f[np.newaxis, :]
29
- assert f.ndim == 2
30
- if normalize:
31
- f /= np.sum(f)
32
- if flip:
33
- f = f[::-1, ::-1]
34
- f = f[:, :, np.newaxis, np.newaxis]
35
- f = np.tile(f, [1, 1, int(x.shape[1]), 1])
36
-
37
- # No-op => early exit.
38
- if f.shape == (1, 1) and f[0,0] == 1:
39
- return x
40
-
41
- # Convolve using depthwise_conv2d.
42
- orig_dtype = x.dtype
43
- x = tf.cast(x, tf.float32) # tf.nn.depthwise_conv2d() doesn't support fp16
44
- f = tf.constant(f, dtype=x.dtype, name='filter')
45
- strides = [1, 1, stride, stride]
46
- x = tf.nn.depthwise_conv2d(x, f, strides=strides, padding='SAME', data_format='NCHW')
47
- x = tf.cast(x, orig_dtype)
48
- return x
49
-
50
- def _upscale2d(x, factor=2, gain=1):
51
- assert x.shape.ndims == 4 and all(dim.value is not None for dim in x.shape[1:])
52
- assert isinstance(factor, int) and factor >= 1
53
-
54
- # Apply gain.
55
- if gain != 1:
56
- x *= gain
57
-
58
- # No-op => early exit.
59
- if factor == 1:
60
- return x
61
-
62
- # Upscale using tf.tile().
63
- s = x.shape
64
- x = tf.reshape(x, [-1, s[1], s[2], 1, s[3], 1])
65
- x = tf.tile(x, [1, 1, 1, factor, 1, factor])
66
- x = tf.reshape(x, [-1, s[1], s[2] * factor, s[3] * factor])
67
- return x
68
-
69
- def _downscale2d(x, factor=2, gain=1):
70
- assert x.shape.ndims == 4 and all(dim.value is not None for dim in x.shape[1:])
71
- assert isinstance(factor, int) and factor >= 1
72
-
73
- # 2x2, float32 => downscale using _blur2d().
74
- if factor == 2 and x.dtype == tf.float32:
75
- f = [np.sqrt(gain) / factor] * factor
76
- return _blur2d(x, f=f, normalize=False, stride=factor)
77
-
78
- # Apply gain.
79
- if gain != 1:
80
- x *= gain
81
-
82
- # No-op => early exit.
83
- if factor == 1:
84
- return x
85
-
86
- # Large factor => downscale using tf.nn.avg_pool().
87
- # NOTE: Requires tf_config['graph_options.place_pruned_graph']=True to work.
88
- ksize = [1, 1, factor, factor]
89
- return tf.nn.avg_pool(x, ksize=ksize, strides=ksize, padding='VALID', data_format='NCHW')
90
-
91
- #----------------------------------------------------------------------------
92
- # High-level ops for manipulating 4D activation tensors.
93
- # The gradients of these are meant to be as efficient as possible.
94
-
95
- def blur2d(x, f=[1,2,1], normalize=True):
96
- with tf.variable_scope('Blur2D'):
97
- @tf.custom_gradient
98
- def func(x):
99
- y = _blur2d(x, f, normalize)
100
- @tf.custom_gradient
101
- def grad(dy):
102
- dx = _blur2d(dy, f, normalize, flip=True)
103
- return dx, lambda ddx: _blur2d(ddx, f, normalize)
104
- return y, grad
105
- return func(x)
106
-
107
- def upscale2d(x, factor=2):
108
- with tf.variable_scope('Upscale2D'):
109
- @tf.custom_gradient
110
- def func(x):
111
- y = _upscale2d(x, factor)
112
- @tf.custom_gradient
113
- def grad(dy):
114
- dx = _downscale2d(dy, factor, gain=factor**2)
115
- return dx, lambda ddx: _upscale2d(ddx, factor)
116
- return y, grad
117
- return func(x)
118
-
119
- def downscale2d(x, factor=2):
120
- with tf.variable_scope('Downscale2D'):
121
- @tf.custom_gradient
122
- def func(x):
123
- y = _downscale2d(x, factor)
124
- @tf.custom_gradient
125
- def grad(dy):
126
- dx = _upscale2d(dy, factor, gain=1/factor**2)
127
- return dx, lambda ddx: _downscale2d(ddx, factor)
128
- return y, grad
129
- return func(x)
130
-
131
- #----------------------------------------------------------------------------
132
- # Get/create weight tensor for a convolutional or fully-connected layer.
133
-
134
- def get_weight(shape, gain=np.sqrt(2), use_wscale=False, lrmul=1):
135
- fan_in = np.prod(shape[:-1]) # [kernel, kernel, fmaps_in, fmaps_out] or [in, out]
136
- he_std = gain / np.sqrt(fan_in) # He init
137
-
138
- # Equalized learning rate and custom learning rate multiplier.
139
- if use_wscale:
140
- init_std = 1.0 / lrmul
141
- runtime_coef = he_std * lrmul
142
- else:
143
- init_std = he_std / lrmul
144
- runtime_coef = lrmul
145
-
146
- # Create variable.
147
- init = tf.initializers.random_normal(0, init_std)
148
- return tf.get_variable('weight', shape=shape, initializer=init) * runtime_coef
149
-
150
- #----------------------------------------------------------------------------
151
- # Fully-connected layer.
152
-
153
- def dense(x, fmaps, **kwargs):
154
- if len(x.shape) > 2:
155
- x = tf.reshape(x, [-1, np.prod([d.value for d in x.shape[1:]])])
156
- w = get_weight([x.shape[1].value, fmaps], **kwargs)
157
- w = tf.cast(w, x.dtype)
158
- return tf.matmul(x, w)
159
-
160
- #----------------------------------------------------------------------------
161
- # Convolutional layer.
162
-
163
- def conv2d(x, fmaps, kernel, **kwargs):
164
- assert kernel >= 1 and kernel % 2 == 1
165
- w = get_weight([kernel, kernel, x.shape[1].value, fmaps], **kwargs)
166
- w = tf.cast(w, x.dtype)
167
- return tf.nn.conv2d(x, w, strides=[1,1,1,1], padding='SAME', data_format='NCHW')
168
-
169
- #----------------------------------------------------------------------------
170
- # Fused convolution + scaling.
171
- # Faster and uses less memory than performing the operations separately.
172
-
173
- def upscale2d_conv2d(x, fmaps, kernel, fused_scale='auto', **kwargs):
174
- assert kernel >= 1 and kernel % 2 == 1
175
- assert fused_scale in [True, False, 'auto']
176
- if fused_scale == 'auto':
177
- fused_scale = min(x.shape[2:]) * 2 >= 128
178
-
179
- # Not fused => call the individual ops directly.
180
- if not fused_scale:
181
- return conv2d(upscale2d(x), fmaps, kernel, **kwargs)
182
-
183
- # Fused => perform both ops simultaneously using tf.nn.conv2d_transpose().
184
- w = get_weight([kernel, kernel, x.shape[1].value, fmaps], **kwargs)
185
- w = tf.transpose(w, [0, 1, 3, 2]) # [kernel, kernel, fmaps_out, fmaps_in]
186
- w = tf.pad(w, [[1,1], [1,1], [0,0], [0,0]], mode='CONSTANT')
187
- w = tf.add_n([w[1:, 1:], w[:-1, 1:], w[1:, :-1], w[:-1, :-1]])
188
- w = tf.cast(w, x.dtype)
189
- os = [tf.shape(x)[0], fmaps, x.shape[2] * 2, x.shape[3] * 2]
190
- return tf.nn.conv2d_transpose(x, w, os, strides=[1,1,2,2], padding='SAME', data_format='NCHW')
191
-
192
- def conv2d_downscale2d(x, fmaps, kernel, fused_scale='auto', **kwargs):
193
- assert kernel >= 1 and kernel % 2 == 1
194
- assert fused_scale in [True, False, 'auto']
195
- if fused_scale == 'auto':
196
- fused_scale = min(x.shape[2:]) >= 128
197
-
198
- # Not fused => call the individual ops directly.
199
- if not fused_scale:
200
- return downscale2d(conv2d(x, fmaps, kernel, **kwargs))
201
-
202
- # Fused => perform both ops simultaneously using tf.nn.conv2d().
203
- w = get_weight([kernel, kernel, x.shape[1].value, fmaps], **kwargs)
204
- w = tf.pad(w, [[1,1], [1,1], [0,0], [0,0]], mode='CONSTANT')
205
- w = tf.add_n([w[1:, 1:], w[:-1, 1:], w[1:, :-1], w[:-1, :-1]]) * 0.25
206
- w = tf.cast(w, x.dtype)
207
- return tf.nn.conv2d(x, w, strides=[1,1,2,2], padding='SAME', data_format='NCHW')
208
-
209
- #----------------------------------------------------------------------------
210
- # Apply bias to the given activation tensor.
211
-
212
- def apply_bias(x, lrmul=1):
213
- b = tf.get_variable('bias', shape=[x.shape[1]], initializer=tf.initializers.zeros()) * lrmul
214
- b = tf.cast(b, x.dtype)
215
- if len(x.shape) == 2:
216
- return x + b
217
- return x + tf.reshape(b, [1, -1, 1, 1])
218
-
219
- #----------------------------------------------------------------------------
220
- # Leaky ReLU activation. More efficient than tf.nn.leaky_relu() and supports FP16.
221
-
222
- def leaky_relu(x, alpha=0.2):
223
- with tf.variable_scope('LeakyReLU'):
224
- alpha = tf.constant(alpha, dtype=x.dtype, name='alpha')
225
- @tf.custom_gradient
226
- def func(x):
227
- y = tf.maximum(x, x * alpha)
228
- @tf.custom_gradient
229
- def grad(dy):
230
- dx = tf.where(y >= 0, dy, dy * alpha)
231
- return dx, lambda ddx: tf.where(y >= 0, ddx, ddx * alpha)
232
- return y, grad
233
- return func(x)
234
-
235
- #----------------------------------------------------------------------------
236
- # Pixelwise feature vector normalization.
237
-
238
- def pixel_norm(x, epsilon=1e-8):
239
- with tf.variable_scope('PixelNorm'):
240
- epsilon = tf.constant(epsilon, dtype=x.dtype, name='epsilon')
241
- return x * tf.rsqrt(tf.reduce_mean(tf.square(x), axis=1, keepdims=True) + epsilon)
242
-
243
- #----------------------------------------------------------------------------
244
- # Instance normalization.
245
-
246
- def instance_norm(x, epsilon=1e-8):
247
- assert len(x.shape) == 4 # NCHW
248
- with tf.variable_scope('InstanceNorm'):
249
- orig_dtype = x.dtype
250
- x = tf.cast(x, tf.float32)
251
- x -= tf.reduce_mean(x, axis=[2,3], keepdims=True)
252
- epsilon = tf.constant(epsilon, dtype=x.dtype, name='epsilon')
253
- x *= tf.rsqrt(tf.reduce_mean(tf.square(x), axis=[2,3], keepdims=True) + epsilon)
254
- x = tf.cast(x, orig_dtype)
255
- return x
256
-
257
- #----------------------------------------------------------------------------
258
- # Style modulation.
259
-
260
- def style_mod(x, dlatent, **kwargs):
261
- with tf.variable_scope('StyleMod'):
262
- style = apply_bias(dense(dlatent, fmaps=x.shape[1]*2, gain=1, **kwargs))
263
- style = tf.reshape(style, [-1, 2, x.shape[1]] + [1] * (len(x.shape) - 2))
264
- return x * (style[:,0] + 1) + style[:,1]
265
-
266
- #----------------------------------------------------------------------------
267
- # Noise input.
268
-
269
- def apply_noise(x, noise_var=None, randomize_noise=True):
270
- assert len(x.shape) == 4 # NCHW
271
- with tf.variable_scope('Noise'):
272
- if noise_var is None or randomize_noise:
273
- noise = tf.random_normal([tf.shape(x)[0], 1, x.shape[2], x.shape[3]], dtype=x.dtype)
274
- else:
275
- noise = tf.cast(noise_var, x.dtype)
276
- weight = tf.get_variable('weight', shape=[x.shape[1].value], initializer=tf.initializers.zeros())
277
- return x + noise * tf.reshape(tf.cast(weight, x.dtype), [1, -1, 1, 1])
278
-
279
- #----------------------------------------------------------------------------
280
- # Minibatch standard deviation.
281
-
282
- def minibatch_stddev_layer(x, group_size=4, num_new_features=1):
283
- with tf.variable_scope('MinibatchStddev'):
284
- group_size = tf.minimum(group_size, tf.shape(x)[0]) # Minibatch must be divisible by (or smaller than) group_size.
285
- s = x.shape # [NCHW] Input shape.
286
- y = tf.reshape(x, [group_size, -1, num_new_features, s[1]//num_new_features, s[2], s[3]]) # [GMncHW] Split minibatch into M groups of size G. Split channels into n channel groups c.
287
- y = tf.cast(y, tf.float32) # [GMncHW] Cast to FP32.
288
- y -= tf.reduce_mean(y, axis=0, keepdims=True) # [GMncHW] Subtract mean over group.
289
- y = tf.reduce_mean(tf.square(y), axis=0) # [MncHW] Calc variance over group.
290
- y = tf.sqrt(y + 1e-8) # [MncHW] Calc stddev over group.
291
- y = tf.reduce_mean(y, axis=[2,3,4], keepdims=True) # [Mn111] Take average over fmaps and pixels.
292
- y = tf.reduce_mean(y, axis=[2]) # [Mn11] Split channels into c channel groups
293
- y = tf.cast(y, x.dtype) # [Mn11] Cast back to original data type.
294
- y = tf.tile(y, [group_size, 1, s[2], s[3]]) # [NnHW] Replicate over group and pixels.
295
- return tf.concat([x, y], axis=1) # [NCHW] Append as new fmap.
296
-
297
- #----------------------------------------------------------------------------
298
- # Style-based generator used in the StyleGAN paper.
299
- # Composed of two sub-networks (G_mapping and G_synthesis) that are defined below.
300
-
301
- def G_style(
302
- latents_in, # First input: Latent vectors (Z) [minibatch, latent_size].
303
- labels_in, # Second input: Conditioning labels [minibatch, label_size].
304
- truncation_psi = 0.7, # Style strength multiplier for the truncation trick. None = disable.
305
- truncation_cutoff = 8, # Number of layers for which to apply the truncation trick. None = disable.
306
- truncation_psi_val = None, # Value for truncation_psi to use during validation.
307
- truncation_cutoff_val = None, # Value for truncation_cutoff to use during validation.
308
- dlatent_avg_beta = 0.995, # Decay for tracking the moving average of W during training. None = disable.
309
- style_mixing_prob = 0.9, # Probability of mixing styles during training. None = disable.
310
- is_training = False, # Network is under training? Enables and disables specific features.
311
- is_validation = False, # Network is under validation? Chooses which value to use for truncation_psi.
312
- is_template_graph = False, # True = template graph constructed by the Network class, False = actual evaluation.
313
- components = dnnlib.EasyDict(), # Container for sub-networks. Retained between calls.
314
- **kwargs): # Arguments for sub-networks (G_mapping and G_synthesis).
315
-
316
- # Validate arguments.
317
- assert not is_training or not is_validation
318
- assert isinstance(components, dnnlib.EasyDict)
319
- if is_validation:
320
- truncation_psi = truncation_psi_val
321
- truncation_cutoff = truncation_cutoff_val
322
- if is_training or (truncation_psi is not None and not tflib.is_tf_expression(truncation_psi) and truncation_psi == 1):
323
- truncation_psi = None
324
- if is_training or (truncation_cutoff is not None and not tflib.is_tf_expression(truncation_cutoff) and truncation_cutoff <= 0):
325
- truncation_cutoff = None
326
- if not is_training or (dlatent_avg_beta is not None and not tflib.is_tf_expression(dlatent_avg_beta) and dlatent_avg_beta == 1):
327
- dlatent_avg_beta = None
328
- if not is_training or (style_mixing_prob is not None and not tflib.is_tf_expression(style_mixing_prob) and style_mixing_prob <= 0):
329
- style_mixing_prob = None
330
-
331
- # Setup components.
332
- if 'synthesis' not in components:
333
- components.synthesis = tflib.Network('G_synthesis', func_name=G_synthesis, **kwargs)
334
- num_layers = components.synthesis.input_shape[1]
335
- dlatent_size = components.synthesis.input_shape[2]
336
- if 'mapping' not in components:
337
- components.mapping = tflib.Network('G_mapping', func_name=G_mapping, dlatent_broadcast=num_layers, **kwargs)
338
-
339
- # Setup variables.
340
- lod_in = tf.get_variable('lod', initializer=np.float32(0), trainable=False)
341
- dlatent_avg = tf.get_variable('dlatent_avg', shape=[dlatent_size], initializer=tf.initializers.zeros(), trainable=False)
342
-
343
- # Evaluate mapping network.
344
- dlatents = components.mapping.get_output_for(latents_in, labels_in, **kwargs)
345
-
346
- # Update moving average of W.
347
- if dlatent_avg_beta is not None:
348
- with tf.variable_scope('DlatentAvg'):
349
- batch_avg = tf.reduce_mean(dlatents[:, 0], axis=0)
350
- update_op = tf.assign(dlatent_avg, tflib.lerp(batch_avg, dlatent_avg, dlatent_avg_beta))
351
- with tf.control_dependencies([update_op]):
352
- dlatents = tf.identity(dlatents)
353
-
354
- # Perform style mixing regularization.
355
- if style_mixing_prob is not None:
356
- with tf.name_scope('StyleMix'):
357
- latents2 = tf.random_normal(tf.shape(latents_in))
358
- dlatents2 = components.mapping.get_output_for(latents2, labels_in, **kwargs)
359
- layer_idx = np.arange(num_layers)[np.newaxis, :, np.newaxis]
360
- cur_layers = num_layers - tf.cast(lod_in, tf.int32) * 2
361
- mixing_cutoff = tf.cond(
362
- tf.random_uniform([], 0.0, 1.0) < style_mixing_prob,
363
- lambda: tf.random_uniform([], 1, cur_layers, dtype=tf.int32),
364
- lambda: cur_layers)
365
- dlatents = tf.where(tf.broadcast_to(layer_idx < mixing_cutoff, tf.shape(dlatents)), dlatents, dlatents2)
366
-
367
- # Apply truncation trick.
368
- if truncation_psi is not None and truncation_cutoff is not None:
369
- with tf.variable_scope('Truncation'):
370
- layer_idx = np.arange(num_layers)[np.newaxis, :, np.newaxis]
371
- ones = np.ones(layer_idx.shape, dtype=np.float32)
372
- coefs = tf.where(layer_idx < truncation_cutoff, truncation_psi * ones, ones)
373
- dlatents = tflib.lerp(dlatent_avg, dlatents, coefs)
374
-
375
- # Evaluate synthesis network.
376
- with tf.control_dependencies([tf.assign(components.synthesis.find_var('lod'), lod_in)]):
377
- images_out = components.synthesis.get_output_for(dlatents, force_clean_graph=is_template_graph, **kwargs)
378
- return tf.identity(images_out, name='images_out')
379
-
380
- #----------------------------------------------------------------------------
381
- # Mapping network used in the StyleGAN paper.
382
-
383
- def G_mapping(
384
- latents_in, # First input: Latent vectors (Z) [minibatch, latent_size].
385
- labels_in, # Second input: Conditioning labels [minibatch, label_size].
386
- latent_size = 512, # Latent vector (Z) dimensionality.
387
- label_size = 0, # Label dimensionality, 0 if no labels.
388
- dlatent_size = 512, # Disentangled latent (W) dimensionality.
389
- dlatent_broadcast = None, # Output disentangled latent (W) as [minibatch, dlatent_size] or [minibatch, dlatent_broadcast, dlatent_size].
390
- mapping_layers = 8, # Number of mapping layers.
391
- mapping_fmaps = 512, # Number of activations in the mapping layers.
392
- mapping_lrmul = 0.01, # Learning rate multiplier for the mapping layers.
393
- mapping_nonlinearity = 'lrelu', # Activation function: 'relu', 'lrelu'.
394
- use_wscale = True, # Enable equalized learning rate?
395
- normalize_latents = True, # Normalize latent vectors (Z) before feeding them to the mapping layers?
396
- dtype = 'float32', # Data type to use for activations and outputs.
397
- **_kwargs): # Ignore unrecognized keyword args.
398
-
399
- act, gain = {'relu': (tf.nn.relu, np.sqrt(2)), 'lrelu': (leaky_relu, np.sqrt(2))}[mapping_nonlinearity]
400
-
401
- # Inputs.
402
- latents_in.set_shape([None, latent_size])
403
- labels_in.set_shape([None, label_size])
404
- latents_in = tf.cast(latents_in, dtype)
405
- labels_in = tf.cast(labels_in, dtype)
406
- x = latents_in
407
-
408
- # Embed labels and concatenate them with latents.
409
- if label_size:
410
- with tf.variable_scope('LabelConcat'):
411
- w = tf.get_variable('weight', shape=[label_size, latent_size], initializer=tf.initializers.random_normal())
412
- y = tf.matmul(labels_in, tf.cast(w, dtype))
413
- x = tf.concat([x, y], axis=1)
414
-
415
- # Normalize latents.
416
- if normalize_latents:
417
- x = pixel_norm(x)
418
-
419
- # Mapping layers.
420
- for layer_idx in range(mapping_layers):
421
- with tf.variable_scope('Dense%d' % layer_idx):
422
- fmaps = dlatent_size if layer_idx == mapping_layers - 1 else mapping_fmaps
423
- x = dense(x, fmaps=fmaps, gain=gain, use_wscale=use_wscale, lrmul=mapping_lrmul)
424
- x = apply_bias(x, lrmul=mapping_lrmul)
425
- x = act(x)
426
-
427
- # Broadcast.
428
- if dlatent_broadcast is not None:
429
- with tf.variable_scope('Broadcast'):
430
- x = tf.tile(x[:, np.newaxis], [1, dlatent_broadcast, 1])
431
-
432
- # Output.
433
- assert x.dtype == tf.as_dtype(dtype)
434
- return tf.identity(x, name='dlatents_out')
435
-
436
- #----------------------------------------------------------------------------
437
- # Synthesis network used in the StyleGAN paper.
438
-
439
- def G_synthesis(
440
- dlatents_in, # Input: Disentangled latents (W) [minibatch, num_layers, dlatent_size].
441
- dlatent_size = 512, # Disentangled latent (W) dimensionality.
442
- num_channels = 3, # Number of output color channels.
443
- resolution = 1024, # Output resolution.
444
- fmap_base = 8192, # Overall multiplier for the number of feature maps.
445
- fmap_decay = 1.0, # log2 feature map reduction when doubling the resolution.
446
- fmap_max = 512, # Maximum number of feature maps in any layer.
447
- use_styles = True, # Enable style inputs?
448
- const_input_layer = True, # First layer is a learned constant?
449
- use_noise = True, # Enable noise inputs?
450
- randomize_noise = True, # True = randomize noise inputs every time (non-deterministic), False = read noise inputs from variables.
451
- nonlinearity = 'lrelu', # Activation function: 'relu', 'lrelu'
452
- use_wscale = True, # Enable equalized learning rate?
453
- use_pixel_norm = False, # Enable pixelwise feature vector normalization?
454
- use_instance_norm = True, # Enable instance normalization?
455
- dtype = 'float32', # Data type to use for activations and outputs.
456
- fused_scale = 'auto', # True = fused convolution + scaling, False = separate ops, 'auto' = decide automatically.
457
- blur_filter = [1,2,1], # Low-pass filter to apply when resampling activations. None = no filtering.
458
- structure = 'auto', # 'fixed' = no progressive growing, 'linear' = human-readable, 'recursive' = efficient, 'auto' = select automatically.
459
- is_template_graph = False, # True = template graph constructed by the Network class, False = actual evaluation.
460
- force_clean_graph = False, # True = construct a clean graph that looks nice in TensorBoard, False = default behavior.
461
- **_kwargs): # Ignore unrecognized keyword args.
462
-
463
- resolution_log2 = int(np.log2(resolution))
464
- assert resolution == 2**resolution_log2 and resolution >= 4
465
- def nf(stage): return min(int(fmap_base / (2.0 ** (stage * fmap_decay))), fmap_max)
466
- def blur(x): return blur2d(x, blur_filter) if blur_filter else x
467
- if is_template_graph: force_clean_graph = True
468
- if force_clean_graph: randomize_noise = False
469
- if structure == 'auto': structure = 'linear' if force_clean_graph else 'recursive'
470
- act, gain = {'relu': (tf.nn.relu, np.sqrt(2)), 'lrelu': (leaky_relu, np.sqrt(2))}[nonlinearity]
471
- num_layers = resolution_log2 * 2 - 2
472
- num_styles = num_layers if use_styles else 1
473
- images_out = None
474
-
475
- # Primary inputs.
476
- dlatents_in.set_shape([None, num_styles, dlatent_size])
477
- dlatents_in = tf.cast(dlatents_in, dtype)
478
- lod_in = tf.cast(tf.get_variable('lod', initializer=np.float32(0), trainable=False), dtype)
479
-
480
- # Noise inputs.
481
- noise_inputs = []
482
- if use_noise:
483
- for layer_idx in range(num_layers):
484
- res = layer_idx // 2 + 2
485
- shape = [1, use_noise, 2**res, 2**res]
486
- noise_inputs.append(tf.get_variable('noise%d' % layer_idx, shape=shape, initializer=tf.initializers.random_normal(), trainable=False))
487
-
488
- # Things to do at the end of each layer.
489
- def layer_epilogue(x, layer_idx):
490
- if use_noise:
491
- x = apply_noise(x, noise_inputs[layer_idx], randomize_noise=randomize_noise)
492
- x = apply_bias(x)
493
- x = act(x)
494
- if use_pixel_norm:
495
- x = pixel_norm(x)
496
- if use_instance_norm:
497
- x = instance_norm(x)
498
- if use_styles:
499
- x = style_mod(x, dlatents_in[:, layer_idx], use_wscale=use_wscale)
500
- return x
501
-
502
- # Early layers.
503
- with tf.variable_scope('4x4'):
504
- if const_input_layer:
505
- with tf.variable_scope('Const'):
506
- x = tf.get_variable('const', shape=[1, nf(1), 4, 4], initializer=tf.initializers.ones())
507
- x = layer_epilogue(tf.tile(tf.cast(x, dtype), [tf.shape(dlatents_in)[0], 1, 1, 1]), 0)
508
- else:
509
- with tf.variable_scope('Dense'):
510
- x = dense(dlatents_in[:, 0], fmaps=nf(1)*16, gain=gain/4, use_wscale=use_wscale) # tweak gain to match the official implementation of Progressing GAN
511
- x = layer_epilogue(tf.reshape(x, [-1, nf(1), 4, 4]), 0)
512
- with tf.variable_scope('Conv'):
513
- x = layer_epilogue(conv2d(x, fmaps=nf(1), kernel=3, gain=gain, use_wscale=use_wscale), 1)
514
-
515
- # Building blocks for remaining layers.
516
- def block(res, x): # res = 3..resolution_log2
517
- with tf.variable_scope('%dx%d' % (2**res, 2**res)):
518
- with tf.variable_scope('Conv0_up'):
519
- x = layer_epilogue(blur(upscale2d_conv2d(x, fmaps=nf(res-1), kernel=3, gain=gain, use_wscale=use_wscale, fused_scale=fused_scale)), res*2-4)
520
- with tf.variable_scope('Conv1'):
521
- x = layer_epilogue(conv2d(x, fmaps=nf(res-1), kernel=3, gain=gain, use_wscale=use_wscale), res*2-3)
522
- return x
523
- def torgb(res, x): # res = 2..resolution_log2
524
- lod = resolution_log2 - res
525
- with tf.variable_scope('ToRGB_lod%d' % lod):
526
- return apply_bias(conv2d(x, fmaps=num_channels, kernel=1, gain=1, use_wscale=use_wscale))
527
-
528
- # Fixed structure: simple and efficient, but does not support progressive growing.
529
- if structure == 'fixed':
530
- for res in range(3, resolution_log2 + 1):
531
- x = block(res, x)
532
- images_out = torgb(resolution_log2, x)
533
-
534
- # Linear structure: simple but inefficient.
535
- if structure == 'linear':
536
- images_out = torgb(2, x)
537
- for res in range(3, resolution_log2 + 1):
538
- lod = resolution_log2 - res
539
- x = block(res, x)
540
- img = torgb(res, x)
541
- images_out = upscale2d(images_out)
542
- with tf.variable_scope('Grow_lod%d' % lod):
543
- images_out = tflib.lerp_clip(img, images_out, lod_in - lod)
544
-
545
- # Recursive structure: complex but efficient.
546
- if structure == 'recursive':
547
- def cset(cur_lambda, new_cond, new_lambda):
548
- return lambda: tf.cond(new_cond, new_lambda, cur_lambda)
549
- def grow(x, res, lod):
550
- y = block(res, x)
551
- img = lambda: upscale2d(torgb(res, y), 2**lod)
552
- img = cset(img, (lod_in > lod), lambda: upscale2d(tflib.lerp(torgb(res, y), upscale2d(torgb(res - 1, x)), lod_in - lod), 2**lod))
553
- if lod > 0: img = cset(img, (lod_in < lod), lambda: grow(y, res + 1, lod - 1))
554
- return img()
555
- images_out = grow(x, 3, resolution_log2 - 3)
556
-
557
- assert images_out.dtype == tf.as_dtype(dtype)
558
- return tf.identity(images_out, name='images_out')
559
-
560
- #----------------------------------------------------------------------------
561
- # Discriminator used in the StyleGAN paper.
562
-
563
- def D_basic(
564
- images_in, # First input: Images [minibatch, channel, height, width].
565
- labels_in, # Second input: Labels [minibatch, label_size].
566
- num_channels = 1, # Number of input color channels. Overridden based on dataset.
567
- resolution = 32, # Input resolution. Overridden based on dataset.
568
- label_size = 0, # Dimensionality of the labels, 0 if no labels. Overridden based on dataset.
569
- fmap_base = 8192, # Overall multiplier for the number of feature maps.
570
- fmap_decay = 1.0, # log2 feature map reduction when doubling the resolution.
571
- fmap_max = 512, # Maximum number of feature maps in any layer.
572
- nonlinearity = 'lrelu', # Activation function: 'relu', 'lrelu',
573
- use_wscale = True, # Enable equalized learning rate?
574
- mbstd_group_size = 4, # Group size for the minibatch standard deviation layer, 0 = disable.
575
- mbstd_num_features = 1, # Number of features for the minibatch standard deviation layer.
576
- dtype = 'float32', # Data type to use for activations and outputs.
577
- fused_scale = 'auto', # True = fused convolution + scaling, False = separate ops, 'auto' = decide automatically.
578
- blur_filter = [1,2,1], # Low-pass filter to apply when resampling activations. None = no filtering.
579
- structure = 'auto', # 'fixed' = no progressive growing, 'linear' = human-readable, 'recursive' = efficient, 'auto' = select automatically.
580
- is_template_graph = False, # True = template graph constructed by the Network class, False = actual evaluation.
581
- **_kwargs): # Ignore unrecognized keyword args.
582
-
583
- resolution_log2 = int(np.log2(resolution))
584
- assert resolution == 2**resolution_log2 and resolution >= 4
585
- def nf(stage): return min(int(fmap_base / (2.0 ** (stage * fmap_decay))), fmap_max)
586
- def blur(x): return blur2d(x, blur_filter) if blur_filter else x
587
- if structure == 'auto': structure = 'linear' if is_template_graph else 'recursive'
588
- act, gain = {'relu': (tf.nn.relu, np.sqrt(2)), 'lrelu': (leaky_relu, np.sqrt(2))}[nonlinearity]
589
-
590
- images_in.set_shape([None, num_channels, resolution, resolution])
591
- labels_in.set_shape([None, label_size])
592
- images_in = tf.cast(images_in, dtype)
593
- labels_in = tf.cast(labels_in, dtype)
594
- lod_in = tf.cast(tf.get_variable('lod', initializer=np.float32(0.0), trainable=False), dtype)
595
- scores_out = None
596
-
597
- # Building blocks.
598
- def fromrgb(x, res): # res = 2..resolution_log2
599
- with tf.variable_scope('FromRGB_lod%d' % (resolution_log2 - res)):
600
- return act(apply_bias(conv2d(x, fmaps=nf(res-1), kernel=1, gain=gain, use_wscale=use_wscale)))
601
- def block(x, res): # res = 2..resolution_log2
602
- with tf.variable_scope('%dx%d' % (2**res, 2**res)):
603
- if res >= 3: # 8x8 and up
604
- with tf.variable_scope('Conv0'):
605
- x = act(apply_bias(conv2d(x, fmaps=nf(res-1), kernel=3, gain=gain, use_wscale=use_wscale)))
606
- with tf.variable_scope('Conv1_down'):
607
- x = act(apply_bias(conv2d_downscale2d(blur(x), fmaps=nf(res-2), kernel=3, gain=gain, use_wscale=use_wscale, fused_scale=fused_scale)))
608
- else: # 4x4
609
- if mbstd_group_size > 1:
610
- x = minibatch_stddev_layer(x, mbstd_group_size, mbstd_num_features)
611
- with tf.variable_scope('Conv'):
612
- x = act(apply_bias(conv2d(x, fmaps=nf(res-1), kernel=3, gain=gain, use_wscale=use_wscale)))
613
- with tf.variable_scope('Dense0'):
614
- x = act(apply_bias(dense(x, fmaps=nf(res-2), gain=gain, use_wscale=use_wscale)))
615
- with tf.variable_scope('Dense1'):
616
- x = apply_bias(dense(x, fmaps=max(label_size, 1), gain=1, use_wscale=use_wscale))
617
- return x
618
-
619
- # Fixed structure: simple and efficient, but does not support progressive growing.
620
- if structure == 'fixed':
621
- x = fromrgb(images_in, resolution_log2)
622
- for res in range(resolution_log2, 2, -1):
623
- x = block(x, res)
624
- scores_out = block(x, 2)
625
-
626
- # Linear structure: simple but inefficient.
627
- if structure == 'linear':
628
- img = images_in
629
- x = fromrgb(img, resolution_log2)
630
- for res in range(resolution_log2, 2, -1):
631
- lod = resolution_log2 - res
632
- x = block(x, res)
633
- img = downscale2d(img)
634
- y = fromrgb(img, res - 1)
635
- with tf.variable_scope('Grow_lod%d' % lod):
636
- x = tflib.lerp_clip(x, y, lod_in - lod)
637
- scores_out = block(x, 2)
638
-
639
- # Recursive structure: complex but efficient.
640
- if structure == 'recursive':
641
- def cset(cur_lambda, new_cond, new_lambda):
642
- return lambda: tf.cond(new_cond, new_lambda, cur_lambda)
643
- def grow(res, lod):
644
- x = lambda: fromrgb(downscale2d(images_in, 2**lod), res)
645
- if lod > 0: x = cset(x, (lod_in < lod), lambda: grow(res + 1, lod - 1))
646
- x = block(x(), res); y = lambda: x
647
- if res > 2: y = cset(y, (lod_in > lod), lambda: tflib.lerp(x, fromrgb(downscale2d(images_in, 2**(lod+1)), res - 1), lod_in - lod))
648
- return y()
649
- scores_out = grow(2, resolution_log2 - 2)
650
-
651
- # Label conditioning from "Which Training Methods for GANs do actually Converge?"
652
- if label_size:
653
- with tf.variable_scope('LabelSwitch'):
654
- scores_out = tf.reduce_sum(scores_out * labels_in, axis=1, keepdims=True)
655
-
656
- assert scores_out.dtype == tf.as_dtype(dtype)
657
- scores_out = tf.identity(scores_out, name='scores_out')
658
- return scores_out
659
-
660
- #----------------------------------------------------------------------------
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
training/networks_stylegan2.py DELETED
@@ -1,697 +0,0 @@
1
- # Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
2
- #
3
- # This work is made available under the Nvidia Source Code License-NC.
4
- # To view a copy of this license, visit
5
- # https://nvlabs.github.io/stylegan2/license.html
6
-
7
- """Network architectures used in the StyleGAN2 paper."""
8
-
9
- import numpy as np
10
- import tensorflow as tf
11
- import dnnlib
12
- import dnnlib.tflib as tflib
13
- from dnnlib.tflib.ops.upfirdn_2d import upsample_2d, downsample_2d, upsample_conv_2d, conv_downsample_2d
14
- from dnnlib.tflib.ops.fused_bias_act import fused_bias_act
15
-
16
- # NOTE: Do not import any application-specific modules here!
17
- # Specify all network parameters as kwargs.
18
-
19
- #----------------------------------------------------------------------------
20
- # Get/create weight tensor for a convolution or fully-connected layer.
21
-
22
- def get_weight(shape, gain=1, use_wscale=True, lrmul=1, weight_var='weight'):
23
- fan_in = np.prod(shape[:-1]) # [kernel, kernel, fmaps_in, fmaps_out] or [in, out]
24
- he_std = gain / np.sqrt(fan_in) # He init
25
-
26
- # Equalized learning rate and custom learning rate multiplier.
27
- if use_wscale:
28
- init_std = 1.0 / lrmul
29
- runtime_coef = he_std * lrmul
30
- else:
31
- init_std = he_std / lrmul
32
- runtime_coef = lrmul
33
-
34
- # Create variable.
35
- init = tf.initializers.random_normal(0, init_std)
36
- return tf.get_variable(weight_var, shape=shape, initializer=init) * runtime_coef
37
-
38
- #----------------------------------------------------------------------------
39
- # Fully-connected layer.
40
-
41
- def dense_layer(x, fmaps, gain=1, use_wscale=True, lrmul=1, weight_var='weight'):
42
- if len(x.shape) > 2:
43
- x = tf.reshape(x, [-1, np.prod([d.value for d in x.shape[1:]])])
44
- w = get_weight([x.shape[1].value, fmaps], gain=gain, use_wscale=use_wscale, lrmul=lrmul, weight_var=weight_var)
45
- w = tf.cast(w, x.dtype)
46
- return tf.matmul(x, w)
47
-
48
- #----------------------------------------------------------------------------
49
- # Convolution layer with optional upsampling or downsampling.
50
-
51
- def conv2d_layer(x, fmaps, kernel, up=False, down=False, resample_kernel=None, gain=1, use_wscale=True, lrmul=1, weight_var='weight'):
52
- assert not (up and down)
53
- assert kernel >= 1 and kernel % 2 == 1
54
- w = get_weight([kernel, kernel, x.shape[1].value, fmaps], gain=gain, use_wscale=use_wscale, lrmul=lrmul, weight_var=weight_var)
55
- if up:
56
- x = upsample_conv_2d(x, tf.cast(w, x.dtype), data_format='NCHW', k=resample_kernel)
57
- elif down:
58
- x = conv_downsample_2d(x, tf.cast(w, x.dtype), data_format='NCHW', k=resample_kernel)
59
- else:
60
- x = tf.nn.conv2d(x, tf.cast(w, x.dtype), data_format='NCHW', strides=[1,1,1,1], padding='SAME')
61
- return x
62
-
63
- #----------------------------------------------------------------------------
64
- # Apply bias and activation func.
65
-
66
- def apply_bias_act(x, act='linear', alpha=None, gain=None, lrmul=1, bias_var='bias'):
67
- b = tf.get_variable(bias_var, shape=[x.shape[1]], initializer=tf.initializers.zeros()) * lrmul
68
- return fused_bias_act(x, b=tf.cast(b, x.dtype), act=act, alpha=alpha, gain=gain)
69
-
70
- #----------------------------------------------------------------------------
71
- # Naive upsampling (nearest neighbor) and downsampling (average pooling).
72
-
73
- def naive_upsample_2d(x, factor=2):
74
- with tf.variable_scope('NaiveUpsample'):
75
- _N, C, H, W = x.shape.as_list()
76
- x = tf.reshape(x, [-1, C, H, 1, W, 1])
77
- x = tf.tile(x, [1, 1, 1, factor, 1, factor])
78
- return tf.reshape(x, [-1, C, H * factor, W * factor])
79
-
80
- def naive_downsample_2d(x, factor=2):
81
- with tf.variable_scope('NaiveDownsample'):
82
- _N, C, H, W = x.shape.as_list()
83
- x = tf.reshape(x, [-1, C, H // factor, factor, W // factor, factor])
84
- return tf.reduce_mean(x, axis=[3,5])
85
-
86
- #----------------------------------------------------------------------------
87
- # Modulated convolution layer.
88
-
89
- def modulated_conv2d_layer(x, y, fmaps, kernel, up=False, down=False, demodulate=True, resample_kernel=None, gain=1, use_wscale=True, lrmul=1, fused_modconv=True, weight_var='weight', mod_weight_var='mod_weight', mod_bias_var='mod_bias'):
90
- assert not (up and down)
91
- assert kernel >= 1 and kernel % 2 == 1
92
-
93
- # Get weight.
94
- w = get_weight([kernel, kernel, x.shape[1].value, fmaps], gain=gain, use_wscale=use_wscale, lrmul=lrmul, weight_var=weight_var)
95
- ww = w[np.newaxis] # [BkkIO] Introduce minibatch dimension.
96
-
97
- # Modulate.
98
- s = dense_layer(y, fmaps=x.shape[1].value, weight_var=mod_weight_var) # [BI] Transform incoming W to style.
99
- s = apply_bias_act(s, bias_var=mod_bias_var) + 1 # [BI] Add bias (initially 1).
100
- ww *= tf.cast(s[:, np.newaxis, np.newaxis, :, np.newaxis], w.dtype) # [BkkIO] Scale input feature maps.
101
-
102
- # Demodulate.
103
- if demodulate:
104
- d = tf.rsqrt(tf.reduce_sum(tf.square(ww), axis=[1,2,3]) + 1e-8) # [BO] Scaling factor.
105
- ww *= d[:, np.newaxis, np.newaxis, np.newaxis, :] # [BkkIO] Scale output feature maps.
106
-
107
- # Reshape/scale input.
108
- if fused_modconv:
109
- x = tf.reshape(x, [1, -1, x.shape[2], x.shape[3]]) # Fused => reshape minibatch to convolution groups.
110
- w = tf.reshape(tf.transpose(ww, [1, 2, 3, 0, 4]), [ww.shape[1], ww.shape[2], ww.shape[3], -1])
111
- else:
112
- x *= tf.cast(s[:, :, np.newaxis, np.newaxis], x.dtype) # [BIhw] Not fused => scale input activations.
113
-
114
- # Convolution with optional up/downsampling.
115
- if up:
116
- x = upsample_conv_2d(x, tf.cast(w, x.dtype), data_format='NCHW', k=resample_kernel)
117
- elif down:
118
- x = conv_downsample_2d(x, tf.cast(w, x.dtype), data_format='NCHW', k=resample_kernel)
119
- else:
120
- x = tf.nn.conv2d(x, tf.cast(w, x.dtype), data_format='NCHW', strides=[1,1,1,1], padding='SAME')
121
-
122
- # Reshape/scale output.
123
- if fused_modconv:
124
- x = tf.reshape(x, [-1, fmaps, x.shape[2], x.shape[3]]) # Fused => reshape convolution groups back to minibatch.
125
- elif demodulate:
126
- x *= tf.cast(d[:, :, np.newaxis, np.newaxis], x.dtype) # [BOhw] Not fused => scale output activations.
127
- return x
128
-
129
- #----------------------------------------------------------------------------
130
- # Minibatch standard deviation layer.
131
-
132
- def minibatch_stddev_layer(x, group_size=4, num_new_features=1):
133
- group_size = tf.minimum(group_size, tf.shape(x)[0]) # Minibatch must be divisible by (or smaller than) group_size.
134
- s = x.shape # [NCHW] Input shape.
135
- y = tf.reshape(x, [group_size, -1, num_new_features, s[1]//num_new_features, s[2], s[3]]) # [GMncHW] Split minibatch into M groups of size G. Split channels into n channel groups c.
136
- y = tf.cast(y, tf.float32) # [GMncHW] Cast to FP32.
137
- y -= tf.reduce_mean(y, axis=0, keepdims=True) # [GMncHW] Subtract mean over group.
138
- y = tf.reduce_mean(tf.square(y), axis=0) # [MncHW] Calc variance over group.
139
- y = tf.sqrt(y + 1e-8) # [MncHW] Calc stddev over group.
140
- y = tf.reduce_mean(y, axis=[2,3,4], keepdims=True) # [Mn111] Take average over fmaps and pixels.
141
- y = tf.reduce_mean(y, axis=[2]) # [Mn11] Split channels into c channel groups
142
- y = tf.cast(y, x.dtype) # [Mn11] Cast back to original data type.
143
- y = tf.tile(y, [group_size, 1, s[2], s[3]]) # [NnHW] Replicate over group and pixels.
144
- return tf.concat([x, y], axis=1) # [NCHW] Append as new fmap.
145
-
146
- #----------------------------------------------------------------------------
147
- # Main generator network.
148
- # Composed of two sub-networks (mapping and synthesis) that are defined below.
149
- # Used in configs B-F (Table 1).
150
-
151
- def G_main(
152
- latents_in, # First input: Latent vectors (Z) [minibatch, latent_size].
153
- labels_in, # Second input: Conditioning labels [minibatch, label_size].
154
- truncation_psi = 0.5, # Style strength multiplier for the truncation trick. None = disable.
155
- truncation_cutoff = None, # Number of layers for which to apply the truncation trick. None = disable.
156
- truncation_psi_val = None, # Value for truncation_psi to use during validation.
157
- truncation_cutoff_val = None, # Value for truncation_cutoff to use during validation.
158
- dlatent_avg_beta = 0.995, # Decay for tracking the moving average of W during training. None = disable.
159
- style_mixing_prob = 0.9, # Probability of mixing styles during training. None = disable.
160
- is_training = False, # Network is under training? Enables and disables specific features.
161
- is_validation = False, # Network is under validation? Chooses which value to use for truncation_psi.
162
- return_dlatents = False, # Return dlatents in addition to the images?
163
- is_template_graph = False, # True = template graph constructed by the Network class, False = actual evaluation.
164
- components = dnnlib.EasyDict(), # Container for sub-networks. Retained between calls.
165
- mapping_func = 'G_mapping', # Build func name for the mapping network.
166
- synthesis_func = 'G_synthesis_stylegan2', # Build func name for the synthesis network.
167
- **kwargs): # Arguments for sub-networks (mapping and synthesis).
168
-
169
- # Validate arguments.
170
- assert not is_training or not is_validation
171
- assert isinstance(components, dnnlib.EasyDict)
172
- if is_validation:
173
- truncation_psi = truncation_psi_val
174
- truncation_cutoff = truncation_cutoff_val
175
- if is_training or (truncation_psi is not None and not tflib.is_tf_expression(truncation_psi) and truncation_psi == 1):
176
- truncation_psi = None
177
- if is_training:
178
- truncation_cutoff = None
179
- if not is_training or (dlatent_avg_beta is not None and not tflib.is_tf_expression(dlatent_avg_beta) and dlatent_avg_beta == 1):
180
- dlatent_avg_beta = None
181
- if not is_training or (style_mixing_prob is not None and not tflib.is_tf_expression(style_mixing_prob) and style_mixing_prob <= 0):
182
- style_mixing_prob = None
183
-
184
- # Setup components.
185
- if 'synthesis' not in components:
186
- components.synthesis = tflib.Network('G_synthesis', func_name=globals()[synthesis_func], **kwargs)
187
- num_layers = components.synthesis.input_shape[1]
188
- dlatent_size = components.synthesis.input_shape[2]
189
- if 'mapping' not in components:
190
- components.mapping = tflib.Network('G_mapping', func_name=globals()[mapping_func], dlatent_broadcast=num_layers, **kwargs)
191
-
192
- # Setup variables.
193
- lod_in = tf.get_variable('lod', initializer=np.float32(0), trainable=False)
194
- dlatent_avg = tf.get_variable('dlatent_avg', shape=[dlatent_size], initializer=tf.initializers.zeros(), trainable=False)
195
-
196
- # Evaluate mapping network.
197
- dlatents = components.mapping.get_output_for(latents_in, labels_in, is_training=is_training, **kwargs)
198
- dlatents = tf.cast(dlatents, tf.float32)
199
-
200
- # Update moving average of W.
201
- if dlatent_avg_beta is not None:
202
- with tf.variable_scope('DlatentAvg'):
203
- batch_avg = tf.reduce_mean(dlatents[:, 0], axis=0)
204
- update_op = tf.assign(dlatent_avg, tflib.lerp(batch_avg, dlatent_avg, dlatent_avg_beta))
205
- with tf.control_dependencies([update_op]):
206
- dlatents = tf.identity(dlatents)
207
-
208
- # Perform style mixing regularization.
209
- if style_mixing_prob is not None:
210
- with tf.variable_scope('StyleMix'):
211
- latents2 = tf.random_normal(tf.shape(latents_in))
212
- dlatents2 = components.mapping.get_output_for(latents2, labels_in, is_training=is_training, **kwargs)
213
- dlatents2 = tf.cast(dlatents2, tf.float32)
214
- layer_idx = np.arange(num_layers)[np.newaxis, :, np.newaxis]
215
- cur_layers = num_layers - tf.cast(lod_in, tf.int32) * 2
216
- mixing_cutoff = tf.cond(
217
- tf.random_uniform([], 0.0, 1.0) < style_mixing_prob,
218
- lambda: tf.random_uniform([], 1, cur_layers, dtype=tf.int32),
219
- lambda: cur_layers)
220
- dlatents = tf.where(tf.broadcast_to(layer_idx < mixing_cutoff, tf.shape(dlatents)), dlatents, dlatents2)
221
-
222
- # Apply truncation trick.
223
- if truncation_psi is not None:
224
- with tf.variable_scope('Truncation'):
225
- layer_idx = np.arange(num_layers)[np.newaxis, :, np.newaxis]
226
- layer_psi = np.ones(layer_idx.shape, dtype=np.float32)
227
- if truncation_cutoff is None:
228
- layer_psi *= truncation_psi
229
- else:
230
- layer_psi = tf.where(layer_idx < truncation_cutoff, layer_psi * truncation_psi, layer_psi)
231
- dlatents = tflib.lerp(dlatent_avg, dlatents, layer_psi)
232
-
233
- # Evaluate synthesis network.
234
- deps = []
235
- if 'lod' in components.synthesis.vars:
236
- deps.append(tf.assign(components.synthesis.vars['lod'], lod_in))
237
- with tf.control_dependencies(deps):
238
- images_out = components.synthesis.get_output_for(dlatents, is_training=is_training, force_clean_graph=is_template_graph, **kwargs)
239
-
240
- # Return requested outputs.
241
- images_out = tf.identity(images_out, name='images_out')
242
- if return_dlatents:
243
- return images_out, dlatents
244
- return images_out
245
-
246
- #----------------------------------------------------------------------------
247
- # Mapping network.
248
- # Transforms the input latent code (z) to the disentangled latent code (w).
249
- # Used in configs B-F (Table 1).
250
-
251
- def G_mapping(
252
- latents_in, # First input: Latent vectors (Z) [minibatch, latent_size].
253
- labels_in, # Second input: Conditioning labels [minibatch, label_size].
254
- latent_size = 512, # Latent vector (Z) dimensionality.
255
- label_size = 0, # Label dimensionality, 0 if no labels.
256
- dlatent_size = 512, # Disentangled latent (W) dimensionality.
257
- dlatent_broadcast = None, # Output disentangled latent (W) as [minibatch, dlatent_size] or [minibatch, dlatent_broadcast, dlatent_size].
258
- mapping_layers = 8, # Number of mapping layers.
259
- mapping_fmaps = 512, # Number of activations in the mapping layers.
260
- mapping_lrmul = 0.01, # Learning rate multiplier for the mapping layers.
261
- mapping_nonlinearity = 'lrelu', # Activation function: 'relu', 'lrelu', etc.
262
- normalize_latents = True, # Normalize latent vectors (Z) before feeding them to the mapping layers?
263
- dtype = 'float32', # Data type to use for activations and outputs.
264
- **_kwargs): # Ignore unrecognized keyword args.
265
-
266
- act = mapping_nonlinearity
267
-
268
- # Inputs.
269
- latents_in.set_shape([None, latent_size])
270
- labels_in.set_shape([None, label_size])
271
- latents_in = tf.cast(latents_in, dtype)
272
- labels_in = tf.cast(labels_in, dtype)
273
- x = latents_in
274
-
275
- # Embed labels and concatenate them with latents.
276
- if label_size:
277
- with tf.variable_scope('LabelConcat'):
278
- w = tf.get_variable('weight', shape=[label_size, latent_size], initializer=tf.initializers.random_normal())
279
- y = tf.matmul(labels_in, tf.cast(w, dtype))
280
- x = tf.concat([x, y], axis=1)
281
-
282
- # Normalize latents.
283
- if normalize_latents:
284
- with tf.variable_scope('Normalize'):
285
- x *= tf.rsqrt(tf.reduce_mean(tf.square(x), axis=1, keepdims=True) + 1e-8)
286
-
287
- # Mapping layers.
288
- for layer_idx in range(mapping_layers):
289
- with tf.variable_scope('Dense%d' % layer_idx):
290
- fmaps = dlatent_size if layer_idx == mapping_layers - 1 else mapping_fmaps
291
- x = apply_bias_act(dense_layer(x, fmaps=fmaps, lrmul=mapping_lrmul), act=act, lrmul=mapping_lrmul)
292
-
293
- # Broadcast.
294
- if dlatent_broadcast is not None:
295
- with tf.variable_scope('Broadcast'):
296
- x = tf.tile(x[:, np.newaxis], [1, dlatent_broadcast, 1])
297
-
298
- # Output.
299
- assert x.dtype == tf.as_dtype(dtype)
300
- return tf.identity(x, name='dlatents_out')
301
-
302
- #----------------------------------------------------------------------------
303
- # StyleGAN synthesis network with revised architecture (Figure 2d).
304
- # Implements progressive growing, but no skip connections or residual nets (Figure 7).
305
- # Used in configs B-D (Table 1).
306
-
307
- def G_synthesis_stylegan_revised(
308
- dlatents_in, # Input: Disentangled latents (W) [minibatch, num_layers, dlatent_size].
309
- dlatent_size = 512, # Disentangled latent (W) dimensionality.
310
- num_channels = 3, # Number of output color channels.
311
- resolution = 1024, # Output resolution.
312
- fmap_base = 16 << 10, # Overall multiplier for the number of feature maps.
313
- fmap_decay = 1.0, # log2 feature map reduction when doubling the resolution.
314
- fmap_min = 1, # Minimum number of feature maps in any layer.
315
- fmap_max = 512, # Maximum number of feature maps in any layer.
316
- randomize_noise = True, # True = randomize noise inputs every time (non-deterministic), False = read noise inputs from variables.
317
- nonlinearity = 'lrelu', # Activation function: 'relu', 'lrelu', etc.
318
- dtype = 'float32', # Data type to use for activations and outputs.
319
- resample_kernel = [1,3,3,1], # Low-pass filter to apply when resampling activations. None = no filtering.
320
- fused_modconv = True, # Implement modulated_conv2d_layer() as a single fused op?
321
- structure = 'auto', # 'fixed' = no progressive growing, 'linear' = human-readable, 'recursive' = efficient, 'auto' = select automatically.
322
- is_template_graph = False, # True = template graph constructed by the Network class, False = actual evaluation.
323
- force_clean_graph = False, # True = construct a clean graph that looks nice in TensorBoard, False = default behavior.
324
- **_kwargs): # Ignore unrecognized keyword args.
325
-
326
- resolution_log2 = int(np.log2(resolution))
327
- assert resolution == 2**resolution_log2 and resolution >= 4
328
- def nf(stage): return np.clip(int(fmap_base / (2.0 ** (stage * fmap_decay))), fmap_min, fmap_max)
329
- if is_template_graph: force_clean_graph = True
330
- if force_clean_graph: randomize_noise = False
331
- if structure == 'auto': structure = 'linear' if force_clean_graph else 'recursive'
332
- act = nonlinearity
333
- num_layers = resolution_log2 * 2 - 2
334
- images_out = None
335
-
336
- # Primary inputs.
337
- dlatents_in.set_shape([None, num_layers, dlatent_size])
338
- dlatents_in = tf.cast(dlatents_in, dtype)
339
- lod_in = tf.cast(tf.get_variable('lod', initializer=np.float32(0), trainable=False), dtype)
340
-
341
- # Noise inputs.
342
- noise_inputs = []
343
- for layer_idx in range(num_layers - 1):
344
- res = (layer_idx + 5) // 2
345
- shape = [1, 1, 2**res, 2**res]
346
- noise_inputs.append(tf.get_variable('noise%d' % layer_idx, shape=shape, initializer=tf.initializers.random_normal(), trainable=False))
347
-
348
- # Single convolution layer with all the bells and whistles.
349
- def layer(x, layer_idx, fmaps, kernel, up=False):
350
- x = modulated_conv2d_layer(x, dlatents_in[:, layer_idx], fmaps=fmaps, kernel=kernel, up=up, resample_kernel=resample_kernel, fused_modconv=fused_modconv)
351
- if randomize_noise:
352
- noise = tf.random_normal([tf.shape(x)[0], 1, x.shape[2], x.shape[3]], dtype=x.dtype)
353
- else:
354
- noise = tf.cast(noise_inputs[layer_idx], x.dtype)
355
- noise_strength = tf.get_variable('noise_strength', shape=[], initializer=tf.initializers.zeros())
356
- x += noise * tf.cast(noise_strength, x.dtype)
357
- return apply_bias_act(x, act=act)
358
-
359
- # Early layers.
360
- with tf.variable_scope('4x4'):
361
- with tf.variable_scope('Const'):
362
- x = tf.get_variable('const', shape=[1, nf(1), 4, 4], initializer=tf.initializers.random_normal())
363
- x = tf.tile(tf.cast(x, dtype), [tf.shape(dlatents_in)[0], 1, 1, 1])
364
- with tf.variable_scope('Conv'):
365
- x = layer(x, layer_idx=0, fmaps=nf(1), kernel=3)
366
-
367
- # Building blocks for remaining layers.
368
- def block(res, x): # res = 3..resolution_log2
369
- with tf.variable_scope('%dx%d' % (2**res, 2**res)):
370
- with tf.variable_scope('Conv0_up'):
371
- x = layer(x, layer_idx=res*2-5, fmaps=nf(res-1), kernel=3, up=True)
372
- with tf.variable_scope('Conv1'):
373
- x = layer(x, layer_idx=res*2-4, fmaps=nf(res-1), kernel=3)
374
- return x
375
- def torgb(res, x): # res = 2..resolution_log2
376
- with tf.variable_scope('ToRGB_lod%d' % (resolution_log2 - res)):
377
- return apply_bias_act(modulated_conv2d_layer(x, dlatents_in[:, res*2-3], fmaps=num_channels, kernel=1, demodulate=False, fused_modconv=fused_modconv))
378
-
379
- # Fixed structure: simple and efficient, but does not support progressive growing.
380
- if structure == 'fixed':
381
- for res in range(3, resolution_log2 + 1):
382
- x = block(res, x)
383
- images_out = torgb(resolution_log2, x)
384
-
385
- # Linear structure: simple but inefficient.
386
- if structure == 'linear':
387
- images_out = torgb(2, x)
388
- for res in range(3, resolution_log2 + 1):
389
- lod = resolution_log2 - res
390
- x = block(res, x)
391
- img = torgb(res, x)
392
- with tf.variable_scope('Upsample_lod%d' % lod):
393
- images_out = upsample_2d(images_out)
394
- with tf.variable_scope('Grow_lod%d' % lod):
395
- images_out = tflib.lerp_clip(img, images_out, lod_in - lod)
396
-
397
- # Recursive structure: complex but efficient.
398
- if structure == 'recursive':
399
- def cset(cur_lambda, new_cond, new_lambda):
400
- return lambda: tf.cond(new_cond, new_lambda, cur_lambda)
401
- def grow(x, res, lod):
402
- y = block(res, x)
403
- img = lambda: naive_upsample_2d(torgb(res, y), factor=2**lod)
404
- img = cset(img, (lod_in > lod), lambda: naive_upsample_2d(tflib.lerp(torgb(res, y), upsample_2d(torgb(res - 1, x)), lod_in - lod), factor=2**lod))
405
- if lod > 0: img = cset(img, (lod_in < lod), lambda: grow(y, res + 1, lod - 1))
406
- return img()
407
- images_out = grow(x, 3, resolution_log2 - 3)
408
-
409
- assert images_out.dtype == tf.as_dtype(dtype)
410
- return tf.identity(images_out, name='images_out')
411
-
412
- #----------------------------------------------------------------------------
413
- # StyleGAN2 synthesis network (Figure 7).
414
- # Implements skip connections and residual nets (Figure 7), but no progressive growing.
415
- # Used in configs E-F (Table 1).
416
-
417
- def G_synthesis_stylegan2(
418
- dlatents_in, # Input: Disentangled latents (W) [minibatch, num_layers, dlatent_size].
419
- dlatent_size = 512, # Disentangled latent (W) dimensionality.
420
- num_channels = 3, # Number of output color channels.
421
- resolution = 1024, # Output resolution.
422
- fmap_base = 16 << 10, # Overall multiplier for the number of feature maps.
423
- fmap_decay = 1.0, # log2 feature map reduction when doubling the resolution.
424
- fmap_min = 1, # Minimum number of feature maps in any layer.
425
- fmap_max = 512, # Maximum number of feature maps in any layer.
426
- randomize_noise = True, # True = randomize noise inputs every time (non-deterministic), False = read noise inputs from variables.
427
- architecture = 'skip', # Architecture: 'orig', 'skip', 'resnet'.
428
- nonlinearity = 'lrelu', # Activation function: 'relu', 'lrelu', etc.
429
- dtype = 'float32', # Data type to use for activations and outputs.
430
- resample_kernel = [1,3,3,1], # Low-pass filter to apply when resampling activations. None = no filtering.
431
- fused_modconv = True, # Implement modulated_conv2d_layer() as a single fused op?
432
- **_kwargs): # Ignore unrecognized keyword args.
433
-
434
- resolution_log2 = int(np.log2(resolution))
435
- assert resolution == 2**resolution_log2 and resolution >= 4
436
- def nf(stage): return np.clip(int(fmap_base / (2.0 ** (stage * fmap_decay))), fmap_min, fmap_max)
437
- assert architecture in ['orig', 'skip', 'resnet']
438
- act = nonlinearity
439
- num_layers = resolution_log2 * 2 - 2
440
- images_out = None
441
-
442
- # Primary inputs.
443
- dlatents_in.set_shape([None, num_layers, dlatent_size])
444
- dlatents_in = tf.cast(dlatents_in, dtype)
445
-
446
- # Noise inputs.
447
- noise_inputs = []
448
- for layer_idx in range(num_layers - 1):
449
- res = (layer_idx + 5) // 2
450
- shape = [1, 1, 2**res, 2**res]
451
- noise_inputs.append(tf.get_variable('noise%d' % layer_idx, shape=shape, initializer=tf.initializers.random_normal(), trainable=False))
452
-
453
- # Single convolution layer with all the bells and whistles.
454
- def layer(x, layer_idx, fmaps, kernel, up=False):
455
- x = modulated_conv2d_layer(x, dlatents_in[:, layer_idx], fmaps=fmaps, kernel=kernel, up=up, resample_kernel=resample_kernel, fused_modconv=fused_modconv)
456
- if randomize_noise:
457
- noise = tf.random_normal([tf.shape(x)[0], 1, x.shape[2], x.shape[3]], dtype=x.dtype)
458
- else:
459
- noise = tf.cast(noise_inputs[layer_idx], x.dtype)
460
- noise_strength = tf.get_variable('noise_strength', shape=[], initializer=tf.initializers.zeros())
461
- x += noise * tf.cast(noise_strength, x.dtype)
462
- return apply_bias_act(x, act=act)
463
-
464
- # Building blocks for main layers.
465
- def block(x, res): # res = 3..resolution_log2
466
- t = x
467
- with tf.variable_scope('Conv0_up'):
468
- x = layer(x, layer_idx=res*2-5, fmaps=nf(res-1), kernel=3, up=True)
469
- with tf.variable_scope('Conv1'):
470
- x = layer(x, layer_idx=res*2-4, fmaps=nf(res-1), kernel=3)
471
- if architecture == 'resnet':
472
- with tf.variable_scope('Skip'):
473
- t = conv2d_layer(t, fmaps=nf(res-1), kernel=1, up=True, resample_kernel=resample_kernel)
474
- x = (x + t) * (1 / np.sqrt(2))
475
- return x
476
- def upsample(y):
477
- with tf.variable_scope('Upsample'):
478
- return upsample_2d(y, k=resample_kernel)
479
- def torgb(x, y, res): # res = 2..resolution_log2
480
- with tf.variable_scope('ToRGB'):
481
- t = apply_bias_act(modulated_conv2d_layer(x, dlatents_in[:, res*2-3], fmaps=num_channels, kernel=1, demodulate=False, fused_modconv=fused_modconv))
482
- return t if y is None else y + t
483
-
484
- # Early layers.
485
- y = None
486
- with tf.variable_scope('4x4'):
487
- with tf.variable_scope('Const'):
488
- x = tf.get_variable('const', shape=[1, nf(1), 4, 4], initializer=tf.initializers.random_normal())
489
- x = tf.tile(tf.cast(x, dtype), [tf.shape(dlatents_in)[0], 1, 1, 1])
490
- with tf.variable_scope('Conv'):
491
- x = layer(x, layer_idx=0, fmaps=nf(1), kernel=3)
492
- if architecture == 'skip':
493
- y = torgb(x, y, 2)
494
-
495
- # Main layers.
496
- for res in range(3, resolution_log2 + 1):
497
- with tf.variable_scope('%dx%d' % (2**res, 2**res)):
498
- x = block(x, res)
499
- if architecture == 'skip':
500
- y = upsample(y)
501
- if architecture == 'skip' or res == resolution_log2:
502
- y = torgb(x, y, res)
503
- images_out = y
504
-
505
- assert images_out.dtype == tf.as_dtype(dtype)
506
- return tf.identity(images_out, name='images_out')
507
-
508
- #----------------------------------------------------------------------------
509
- # Original StyleGAN discriminator.
510
- # Used in configs B-D (Table 1).
511
-
512
- def D_stylegan(
513
- images_in, # First input: Images [minibatch, channel, height, width].
514
- labels_in, # Second input: Labels [minibatch, label_size].
515
- num_channels = 3, # Number of input color channels. Overridden based on dataset.
516
- resolution = 1024, # Input resolution. Overridden based on dataset.
517
- label_size = 0, # Dimensionality of the labels, 0 if no labels. Overridden based on dataset.
518
- fmap_base = 16 << 10, # Overall multiplier for the number of feature maps.
519
- fmap_decay = 1.0, # log2 feature map reduction when doubling the resolution.
520
- fmap_min = 1, # Minimum number of feature maps in any layer.
521
- fmap_max = 512, # Maximum number of feature maps in any layer.
522
- nonlinearity = 'lrelu', # Activation function: 'relu', 'lrelu', etc.
523
- mbstd_group_size = 4, # Group size for the minibatch standard deviation layer, 0 = disable.
524
- mbstd_num_features = 1, # Number of features for the minibatch standard deviation layer.
525
- dtype = 'float32', # Data type to use for activations and outputs.
526
- resample_kernel = [1,3,3,1], # Low-pass filter to apply when resampling activations. None = no filtering.
527
- structure = 'auto', # 'fixed' = no progressive growing, 'linear' = human-readable, 'recursive' = efficient, 'auto' = select automatically.
528
- is_template_graph = False, # True = template graph constructed by the Network class, False = actual evaluation.
529
- **_kwargs): # Ignore unrecognized keyword args.
530
-
531
- resolution_log2 = int(np.log2(resolution))
532
- assert resolution == 2**resolution_log2 and resolution >= 4
533
- def nf(stage): return np.clip(int(fmap_base / (2.0 ** (stage * fmap_decay))), fmap_min, fmap_max)
534
- if structure == 'auto': structure = 'linear' if is_template_graph else 'recursive'
535
- act = nonlinearity
536
-
537
- images_in.set_shape([None, num_channels, resolution, resolution])
538
- labels_in.set_shape([None, label_size])
539
- images_in = tf.cast(images_in, dtype)
540
- labels_in = tf.cast(labels_in, dtype)
541
- lod_in = tf.cast(tf.get_variable('lod', initializer=np.float32(0.0), trainable=False), dtype)
542
-
543
- # Building blocks for spatial layers.
544
- def fromrgb(x, res): # res = 2..resolution_log2
545
- with tf.variable_scope('FromRGB_lod%d' % (resolution_log2 - res)):
546
- return apply_bias_act(conv2d_layer(x, fmaps=nf(res-1), kernel=1), act=act)
547
- def block(x, res): # res = 2..resolution_log2
548
- with tf.variable_scope('%dx%d' % (2**res, 2**res)):
549
- with tf.variable_scope('Conv0'):
550
- x = apply_bias_act(conv2d_layer(x, fmaps=nf(res-1), kernel=3), act=act)
551
- with tf.variable_scope('Conv1_down'):
552
- x = apply_bias_act(conv2d_layer(x, fmaps=nf(res-2), kernel=3, down=True, resample_kernel=resample_kernel), act=act)
553
- return x
554
-
555
- # Fixed structure: simple and efficient, but does not support progressive growing.
556
- if structure == 'fixed':
557
- x = fromrgb(images_in, resolution_log2)
558
- for res in range(resolution_log2, 2, -1):
559
- x = block(x, res)
560
-
561
- # Linear structure: simple but inefficient.
562
- if structure == 'linear':
563
- img = images_in
564
- x = fromrgb(img, resolution_log2)
565
- for res in range(resolution_log2, 2, -1):
566
- lod = resolution_log2 - res
567
- x = block(x, res)
568
- with tf.variable_scope('Downsample_lod%d' % lod):
569
- img = downsample_2d(img)
570
- y = fromrgb(img, res - 1)
571
- with tf.variable_scope('Grow_lod%d' % lod):
572
- x = tflib.lerp_clip(x, y, lod_in - lod)
573
-
574
- # Recursive structure: complex but efficient.
575
- if structure == 'recursive':
576
- def cset(cur_lambda, new_cond, new_lambda):
577
- return lambda: tf.cond(new_cond, new_lambda, cur_lambda)
578
- def grow(res, lod):
579
- x = lambda: fromrgb(naive_downsample_2d(images_in, factor=2**lod), res)
580
- if lod > 0: x = cset(x, (lod_in < lod), lambda: grow(res + 1, lod - 1))
581
- x = block(x(), res); y = lambda: x
582
- y = cset(y, (lod_in > lod), lambda: tflib.lerp(x, fromrgb(naive_downsample_2d(images_in, factor=2**(lod+1)), res - 1), lod_in - lod))
583
- return y()
584
- x = grow(3, resolution_log2 - 3)
585
-
586
- # Final layers at 4x4 resolution.
587
- with tf.variable_scope('4x4'):
588
- if mbstd_group_size > 1:
589
- with tf.variable_scope('MinibatchStddev'):
590
- x = minibatch_stddev_layer(x, mbstd_group_size, mbstd_num_features)
591
- with tf.variable_scope('Conv'):
592
- x = apply_bias_act(conv2d_layer(x, fmaps=nf(1), kernel=3), act=act)
593
- with tf.variable_scope('Dense0'):
594
- x = apply_bias_act(dense_layer(x, fmaps=nf(0)), act=act)
595
-
596
- # Output layer with label conditioning from "Which Training Methods for GANs do actually Converge?"
597
- with tf.variable_scope('Output'):
598
- x = apply_bias_act(dense_layer(x, fmaps=max(labels_in.shape[1], 1)))
599
- if labels_in.shape[1] > 0:
600
- x = tf.reduce_sum(x * labels_in, axis=1, keepdims=True)
601
- scores_out = x
602
-
603
- # Output.
604
- assert scores_out.dtype == tf.as_dtype(dtype)
605
- scores_out = tf.identity(scores_out, name='scores_out')
606
- return scores_out
607
-
608
- #----------------------------------------------------------------------------
609
- # StyleGAN2 discriminator (Figure 7).
610
- # Implements skip connections and residual nets (Figure 7), but no progressive growing.
611
- # Used in configs E-F (Table 1).
612
-
613
- def D_stylegan2(
614
- images_in, # First input: Images [minibatch, channel, height, width].
615
- labels_in, # Second input: Labels [minibatch, label_size].
616
- num_channels = 3, # Number of input color channels. Overridden based on dataset.
617
- resolution = 1024, # Input resolution. Overridden based on dataset.
618
- label_size = 0, # Dimensionality of the labels, 0 if no labels. Overridden based on dataset.
619
- fmap_base = 16 << 10, # Overall multiplier for the number of feature maps.
620
- fmap_decay = 1.0, # log2 feature map reduction when doubling the resolution.
621
- fmap_min = 1, # Minimum number of feature maps in any layer.
622
- fmap_max = 512, # Maximum number of feature maps in any layer.
623
- architecture = 'resnet', # Architecture: 'orig', 'skip', 'resnet'.
624
- nonlinearity = 'lrelu', # Activation function: 'relu', 'lrelu', etc.
625
- mbstd_group_size = 4, # Group size for the minibatch standard deviation layer, 0 = disable.
626
- mbstd_num_features = 1, # Number of features for the minibatch standard deviation layer.
627
- dtype = 'float32', # Data type to use for activations and outputs.
628
- resample_kernel = [1,3,3,1], # Low-pass filter to apply when resampling activations. None = no filtering.
629
- **_kwargs): # Ignore unrecognized keyword args.
630
-
631
- resolution_log2 = int(np.log2(resolution))
632
- assert resolution == 2**resolution_log2 and resolution >= 4
633
- def nf(stage): return np.clip(int(fmap_base / (2.0 ** (stage * fmap_decay))), fmap_min, fmap_max)
634
- assert architecture in ['orig', 'skip', 'resnet']
635
- act = nonlinearity
636
-
637
- images_in.set_shape([None, num_channels, resolution, resolution])
638
- labels_in.set_shape([None, label_size])
639
- images_in = tf.cast(images_in, dtype)
640
- labels_in = tf.cast(labels_in, dtype)
641
-
642
- # Building blocks for main layers.
643
- def fromrgb(x, y, res): # res = 2..resolution_log2
644
- with tf.variable_scope('FromRGB'):
645
- t = apply_bias_act(conv2d_layer(y, fmaps=nf(res-1), kernel=1), act=act)
646
- return t if x is None else x + t
647
- def block(x, res): # res = 2..resolution_log2
648
- t = x
649
- with tf.variable_scope('Conv0'):
650
- x = apply_bias_act(conv2d_layer(x, fmaps=nf(res-1), kernel=3), act=act)
651
- with tf.variable_scope('Conv1_down'):
652
- x = apply_bias_act(conv2d_layer(x, fmaps=nf(res-2), kernel=3, down=True, resample_kernel=resample_kernel), act=act)
653
- if architecture == 'resnet':
654
- with tf.variable_scope('Skip'):
655
- t = conv2d_layer(t, fmaps=nf(res-2), kernel=1, down=True, resample_kernel=resample_kernel)
656
- x = (x + t) * (1 / np.sqrt(2))
657
- return x
658
- def downsample(y):
659
- with tf.variable_scope('Downsample'):
660
- return downsample_2d(y, k=resample_kernel)
661
-
662
- # Main layers.
663
- x = None
664
- y = images_in
665
- for res in range(resolution_log2, 2, -1):
666
- with tf.variable_scope('%dx%d' % (2**res, 2**res)):
667
- if architecture == 'skip' or res == resolution_log2:
668
- x = fromrgb(x, y, res)
669
- x = block(x, res)
670
- if architecture == 'skip':
671
- y = downsample(y)
672
-
673
- # Final layers.
674
- with tf.variable_scope('4x4'):
675
- if architecture == 'skip':
676
- x = fromrgb(x, y, 2)
677
- if mbstd_group_size > 1:
678
- with tf.variable_scope('MinibatchStddev'):
679
- x = minibatch_stddev_layer(x, mbstd_group_size, mbstd_num_features)
680
- with tf.variable_scope('Conv'):
681
- x = apply_bias_act(conv2d_layer(x, fmaps=nf(1), kernel=3), act=act)
682
- with tf.variable_scope('Dense0'):
683
- x = apply_bias_act(dense_layer(x, fmaps=nf(0)), act=act)
684
-
685
- # Output layer with label conditioning from "Which Training Methods for GANs do actually Converge?"
686
- with tf.variable_scope('Output'):
687
- x = apply_bias_act(dense_layer(x, fmaps=max(labels_in.shape[1], 1)))
688
- if labels_in.shape[1] > 0:
689
- x = tf.reduce_sum(x * labels_in, axis=1, keepdims=True)
690
- scores_out = x
691
-
692
- # Output.
693
- assert scores_out.dtype == tf.as_dtype(dtype)
694
- scores_out = tf.identity(scores_out, name='scores_out')
695
- return scores_out
696
-
697
- #----------------------------------------------------------------------------
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
training/training_loop.py DELETED
@@ -1,356 +0,0 @@
1
- # Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
2
- #
3
- # This work is made available under the Nvidia Source Code License-NC.
4
- # To view a copy of this license, visit
5
- # https://nvlabs.github.io/stylegan2/license.html
6
-
7
- """Main training script."""
8
-
9
- import numpy as np
10
- import tensorflow as tf
11
- import dnnlib
12
- import dnnlib.tflib as tflib
13
- from dnnlib.tflib.autosummary import autosummary
14
-
15
- from training import dataset
16
- from training import misc
17
- from metrics import metric_base
18
-
19
- #----------------------------------------------------------------------------
20
- # Just-in-time processing of training images before feeding them to the networks.
21
-
22
- def process_reals(x, labels, lod, mirror_augment, drange_data, drange_net):
23
- with tf.name_scope('DynamicRange'):
24
- x = tf.cast(x, tf.float32)
25
- x = misc.adjust_dynamic_range(x, drange_data, drange_net)
26
- if mirror_augment:
27
- with tf.name_scope('MirrorAugment'):
28
- x = tf.where(tf.random_uniform([tf.shape(x)[0]]) < 0.5, x, tf.reverse(x, [3]))
29
- with tf.name_scope('FadeLOD'): # Smooth crossfade between consecutive levels-of-detail.
30
- s = tf.shape(x)
31
- y = tf.reshape(x, [-1, s[1], s[2]//2, 2, s[3]//2, 2])
32
- y = tf.reduce_mean(y, axis=[3, 5], keepdims=True)
33
- y = tf.tile(y, [1, 1, 1, 2, 1, 2])
34
- y = tf.reshape(y, [-1, s[1], s[2], s[3]])
35
- x = tflib.lerp(x, y, lod - tf.floor(lod))
36
- with tf.name_scope('UpscaleLOD'): # Upscale to match the expected input/output size of the networks.
37
- s = tf.shape(x)
38
- factor = tf.cast(2 ** tf.floor(lod), tf.int32)
39
- x = tf.reshape(x, [-1, s[1], s[2], 1, s[3], 1])
40
- x = tf.tile(x, [1, 1, 1, factor, 1, factor])
41
- x = tf.reshape(x, [-1, s[1], s[2] * factor, s[3] * factor])
42
- return x, labels
43
-
44
- #----------------------------------------------------------------------------
45
- # Evaluate time-varying training parameters.
46
-
47
- def training_schedule(
48
- cur_nimg,
49
- training_set,
50
- lod_initial_resolution = None, # Image resolution used at the beginning.
51
- lod_training_kimg = 600, # Thousands of real images to show before doubling the resolution.
52
- lod_transition_kimg = 600, # Thousands of real images to show when fading in new layers.
53
- minibatch_size_base = 32, # Global minibatch size.
54
- minibatch_size_dict = {}, # Resolution-specific overrides.
55
- minibatch_gpu_base = 4, # Number of samples processed at a time by one GPU.
56
- minibatch_gpu_dict = {}, # Resolution-specific overrides.
57
- G_lrate_base = 0.002, # Learning rate for the generator.
58
- G_lrate_dict = {}, # Resolution-specific overrides.
59
- D_lrate_base = 0.002, # Learning rate for the discriminator.
60
- D_lrate_dict = {}, # Resolution-specific overrides.
61
- lrate_rampup_kimg = 0, # Duration of learning rate ramp-up.
62
- tick_kimg_base = 4, # Default interval of progress snapshots.
63
- tick_kimg_dict = {8:28, 16:24, 32:20, 64:16, 128:12, 256:8, 512:6, 1024:4}): # Resolution-specific overrides.
64
-
65
- # Initialize result dict.
66
- s = dnnlib.EasyDict()
67
- s.kimg = cur_nimg / 1000.0
68
-
69
- # Training phase.
70
- phase_dur = lod_training_kimg + lod_transition_kimg
71
- phase_idx = int(np.floor(s.kimg / phase_dur)) if phase_dur > 0 else 0
72
- phase_kimg = s.kimg - phase_idx * phase_dur
73
-
74
- # Level-of-detail and resolution.
75
- if lod_initial_resolution is None:
76
- s.lod = 0.0
77
- else:
78
- s.lod = training_set.resolution_log2
79
- s.lod -= np.floor(np.log2(lod_initial_resolution))
80
- s.lod -= phase_idx
81
- if lod_transition_kimg > 0:
82
- s.lod -= max(phase_kimg - lod_training_kimg, 0.0) / lod_transition_kimg
83
- s.lod = max(s.lod, 0.0)
84
- s.resolution = 2 ** (training_set.resolution_log2 - int(np.floor(s.lod)))
85
-
86
- # Minibatch size.
87
- s.minibatch_size = minibatch_size_dict.get(s.resolution, minibatch_size_base)
88
- s.minibatch_gpu = minibatch_gpu_dict.get(s.resolution, minibatch_gpu_base)
89
-
90
- # Learning rate.
91
- s.G_lrate = G_lrate_dict.get(s.resolution, G_lrate_base)
92
- s.D_lrate = D_lrate_dict.get(s.resolution, D_lrate_base)
93
- if lrate_rampup_kimg > 0:
94
- rampup = min(s.kimg / lrate_rampup_kimg, 1.0)
95
- s.G_lrate *= rampup
96
- s.D_lrate *= rampup
97
-
98
- # Other parameters.
99
- s.tick_kimg = tick_kimg_dict.get(s.resolution, tick_kimg_base)
100
- return s
101
-
102
- #----------------------------------------------------------------------------
103
- # Main training script.
104
-
105
- def training_loop(
106
- G_args = {}, # Options for generator network.
107
- D_args = {}, # Options for discriminator network.
108
- G_opt_args = {}, # Options for generator optimizer.
109
- D_opt_args = {}, # Options for discriminator optimizer.
110
- G_loss_args = {}, # Options for generator loss.
111
- D_loss_args = {}, # Options for discriminator loss.
112
- dataset_args = {}, # Options for dataset.load_dataset().
113
- sched_args = {}, # Options for train.TrainingSchedule.
114
- grid_args = {}, # Options for train.setup_snapshot_image_grid().
115
- metric_arg_list = [], # Options for MetricGroup.
116
- tf_config = {}, # Options for tflib.init_tf().
117
- data_dir = None, # Directory to load datasets from.
118
- G_smoothing_kimg = 10.0, # Half-life of the running average of generator weights.
119
- minibatch_repeats = 4, # Number of minibatches to run before adjusting training parameters.
120
- lazy_regularization = True, # Perform regularization as a separate training step?
121
- G_reg_interval = 4, # How often the perform regularization for G? Ignored if lazy_regularization=False.
122
- D_reg_interval = 16, # How often the perform regularization for D? Ignored if lazy_regularization=False.
123
- reset_opt_for_new_lod = True, # Reset optimizer internal state (e.g. Adam moments) when new layers are introduced?
124
- total_kimg = 25000, # Total length of the training, measured in thousands of real images.
125
- mirror_augment = False, # Enable mirror augment?
126
- drange_net = [-1,1], # Dynamic range used when feeding image data to the networks.
127
- image_snapshot_ticks = 50, # How often to save image snapshots? None = only save 'reals.png' and 'fakes-init.png'.
128
- network_snapshot_ticks = 50, # How often to save network snapshots? None = only save 'networks-final.pkl'.
129
- save_tf_graph = False, # Include full TensorFlow computation graph in the tfevents file?
130
- save_weight_histograms = False, # Include weight histograms in the tfevents file?
131
- resume_pkl = None, # Network pickle to resume training from, None = train from scratch.
132
- resume_kimg = 0.0, # Assumed training progress at the beginning. Affects reporting and training schedule.
133
- resume_time = 0.0, # Assumed wallclock time at the beginning. Affects reporting.
134
- resume_with_new_nets = False): # Construct new networks according to G_args and D_args before resuming training?
135
-
136
- # Initialize dnnlib and TensorFlow.
137
- tflib.init_tf(tf_config)
138
- num_gpus = dnnlib.submit_config.num_gpus
139
-
140
- # Load training set.
141
- training_set = dataset.load_dataset(data_dir=dnnlib.convert_path(data_dir), verbose=True, **dataset_args)
142
- grid_size, grid_reals, grid_labels = misc.setup_snapshot_image_grid(training_set, **grid_args)
143
- misc.save_image_grid(grid_reals, dnnlib.make_run_dir_path('reals.png'), drange=training_set.dynamic_range, grid_size=grid_size)
144
-
145
- # Construct or load networks.
146
- with tf.device('/gpu:0'):
147
- if resume_pkl is None or resume_with_new_nets:
148
- print('Constructing networks...')
149
- G = tflib.Network('G', num_channels=training_set.shape[0], resolution=training_set.shape[1], label_size=training_set.label_size, **G_args)
150
- D = tflib.Network('D', num_channels=training_set.shape[0], resolution=training_set.shape[1], label_size=training_set.label_size, **D_args)
151
- Gs = G.clone('Gs')
152
- if resume_pkl is not None:
153
- print('Loading networks from "%s"...' % resume_pkl)
154
- rG, rD, rGs = misc.load_pkl(resume_pkl)
155
- if resume_with_new_nets: G.copy_vars_from(rG); D.copy_vars_from(rD); Gs.copy_vars_from(rGs)
156
- else: G = rG; D = rD; Gs = rGs
157
-
158
- # Print layers and generate initial image snapshot.
159
- G.print_layers(); D.print_layers()
160
- sched = training_schedule(cur_nimg=total_kimg*1000, training_set=training_set, **sched_args)
161
- grid_latents = np.random.randn(np.prod(grid_size), *G.input_shape[1:])
162
- grid_fakes = Gs.run(grid_latents, grid_labels, is_validation=True, minibatch_size=sched.minibatch_gpu)
163
- misc.save_image_grid(grid_fakes, dnnlib.make_run_dir_path('fakes_init.png'), drange=drange_net, grid_size=grid_size)
164
-
165
- # Setup training inputs.
166
- print('Building TensorFlow graph...')
167
- with tf.name_scope('Inputs'), tf.device('/cpu:0'):
168
- lod_in = tf.placeholder(tf.float32, name='lod_in', shape=[])
169
- lrate_in = tf.placeholder(tf.float32, name='lrate_in', shape=[])
170
- minibatch_size_in = tf.placeholder(tf.int32, name='minibatch_size_in', shape=[])
171
- minibatch_gpu_in = tf.placeholder(tf.int32, name='minibatch_gpu_in', shape=[])
172
- minibatch_multiplier = minibatch_size_in // (minibatch_gpu_in * num_gpus)
173
- Gs_beta = 0.5 ** tf.div(tf.cast(minibatch_size_in, tf.float32), G_smoothing_kimg * 1000.0) if G_smoothing_kimg > 0.0 else 0.0
174
-
175
- # Setup optimizers.
176
- G_opt_args = dict(G_opt_args)
177
- D_opt_args = dict(D_opt_args)
178
- for args, reg_interval in [(G_opt_args, G_reg_interval), (D_opt_args, D_reg_interval)]:
179
- args['minibatch_multiplier'] = minibatch_multiplier
180
- args['learning_rate'] = lrate_in
181
- if lazy_regularization:
182
- mb_ratio = reg_interval / (reg_interval + 1)
183
- args['learning_rate'] *= mb_ratio
184
- if 'beta1' in args: args['beta1'] **= mb_ratio
185
- if 'beta2' in args: args['beta2'] **= mb_ratio
186
- G_opt = tflib.Optimizer(name='TrainG', **G_opt_args)
187
- D_opt = tflib.Optimizer(name='TrainD', **D_opt_args)
188
- G_reg_opt = tflib.Optimizer(name='RegG', share=G_opt, **G_opt_args)
189
- D_reg_opt = tflib.Optimizer(name='RegD', share=D_opt, **D_opt_args)
190
-
191
- # Build training graph for each GPU.
192
- data_fetch_ops = []
193
- for gpu in range(num_gpus):
194
- with tf.name_scope('GPU%d' % gpu), tf.device('/gpu:%d' % gpu):
195
-
196
- # Create GPU-specific shadow copies of G and D.
197
- G_gpu = G if gpu == 0 else G.clone(G.name + '_shadow')
198
- D_gpu = D if gpu == 0 else D.clone(D.name + '_shadow')
199
-
200
- # Fetch training data via temporary variables.
201
- with tf.name_scope('DataFetch'):
202
- sched = training_schedule(cur_nimg=int(resume_kimg*1000), training_set=training_set, **sched_args)
203
- reals_var = tf.Variable(name='reals', trainable=False, initial_value=tf.zeros([sched.minibatch_gpu] + training_set.shape))
204
- labels_var = tf.Variable(name='labels', trainable=False, initial_value=tf.zeros([sched.minibatch_gpu, training_set.label_size]))
205
- reals_write, labels_write = training_set.get_minibatch_tf()
206
- reals_write, labels_write = process_reals(reals_write, labels_write, lod_in, mirror_augment, training_set.dynamic_range, drange_net)
207
- reals_write = tf.concat([reals_write, reals_var[minibatch_gpu_in:]], axis=0)
208
- labels_write = tf.concat([labels_write, labels_var[minibatch_gpu_in:]], axis=0)
209
- data_fetch_ops += [tf.assign(reals_var, reals_write)]
210
- data_fetch_ops += [tf.assign(labels_var, labels_write)]
211
- reals_read = reals_var[:minibatch_gpu_in]
212
- labels_read = labels_var[:minibatch_gpu_in]
213
-
214
- # Evaluate loss functions.
215
- lod_assign_ops = []
216
- if 'lod' in G_gpu.vars: lod_assign_ops += [tf.assign(G_gpu.vars['lod'], lod_in)]
217
- if 'lod' in D_gpu.vars: lod_assign_ops += [tf.assign(D_gpu.vars['lod'], lod_in)]
218
- with tf.control_dependencies(lod_assign_ops):
219
- with tf.name_scope('G_loss'):
220
- G_loss, G_reg = dnnlib.util.call_func_by_name(G=G_gpu, D=D_gpu, opt=G_opt, training_set=training_set, minibatch_size=minibatch_gpu_in, **G_loss_args)
221
- with tf.name_scope('D_loss'):
222
- D_loss, D_reg = dnnlib.util.call_func_by_name(G=G_gpu, D=D_gpu, opt=D_opt, training_set=training_set, minibatch_size=minibatch_gpu_in, reals=reals_read, labels=labels_read, **D_loss_args)
223
-
224
- # Register gradients.
225
- if not lazy_regularization:
226
- if G_reg is not None: G_loss += G_reg
227
- if D_reg is not None: D_loss += D_reg
228
- else:
229
- if G_reg is not None: G_reg_opt.register_gradients(tf.reduce_mean(G_reg * G_reg_interval), G_gpu.trainables)
230
- if D_reg is not None: D_reg_opt.register_gradients(tf.reduce_mean(D_reg * D_reg_interval), D_gpu.trainables)
231
- G_opt.register_gradients(tf.reduce_mean(G_loss), G_gpu.trainables)
232
- D_opt.register_gradients(tf.reduce_mean(D_loss), D_gpu.trainables)
233
-
234
- # Setup training ops.
235
- data_fetch_op = tf.group(*data_fetch_ops)
236
- G_train_op = G_opt.apply_updates()
237
- D_train_op = D_opt.apply_updates()
238
- G_reg_op = G_reg_opt.apply_updates(allow_no_op=True)
239
- D_reg_op = D_reg_opt.apply_updates(allow_no_op=True)
240
- Gs_update_op = Gs.setup_as_moving_average_of(G, beta=Gs_beta)
241
-
242
- # Finalize graph.
243
- with tf.device('/gpu:0'):
244
- try:
245
- peak_gpu_mem_op = tf.contrib.memory_stats.MaxBytesInUse()
246
- except tf.errors.NotFoundError:
247
- peak_gpu_mem_op = tf.constant(0)
248
- tflib.init_uninitialized_vars()
249
-
250
- print('Initializing logs...')
251
- summary_log = tf.summary.FileWriter(dnnlib.make_run_dir_path())
252
- if save_tf_graph:
253
- summary_log.add_graph(tf.get_default_graph())
254
- if save_weight_histograms:
255
- G.setup_weight_histograms(); D.setup_weight_histograms()
256
- metrics = metric_base.MetricGroup(metric_arg_list)
257
-
258
- print('Training for %d kimg...\n' % total_kimg)
259
- dnnlib.RunContext.get().update('', cur_epoch=resume_kimg, max_epoch=total_kimg)
260
- maintenance_time = dnnlib.RunContext.get().get_last_update_interval()
261
- cur_nimg = int(resume_kimg * 1000)
262
- cur_tick = -1
263
- tick_start_nimg = cur_nimg
264
- prev_lod = -1.0
265
- running_mb_counter = 0
266
- while cur_nimg < total_kimg * 1000:
267
- if dnnlib.RunContext.get().should_stop(): break
268
-
269
- # Choose training parameters and configure training ops.
270
- sched = training_schedule(cur_nimg=cur_nimg, training_set=training_set, **sched_args)
271
- assert sched.minibatch_size % (sched.minibatch_gpu * num_gpus) == 0
272
- training_set.configure(sched.minibatch_gpu, sched.lod)
273
- if reset_opt_for_new_lod:
274
- if np.floor(sched.lod) != np.floor(prev_lod) or np.ceil(sched.lod) != np.ceil(prev_lod):
275
- G_opt.reset_optimizer_state(); D_opt.reset_optimizer_state()
276
- prev_lod = sched.lod
277
-
278
- # Run training ops.
279
- feed_dict = {lod_in: sched.lod, lrate_in: sched.G_lrate, minibatch_size_in: sched.minibatch_size, minibatch_gpu_in: sched.minibatch_gpu}
280
- for _repeat in range(minibatch_repeats):
281
- rounds = range(0, sched.minibatch_size, sched.minibatch_gpu * num_gpus)
282
- run_G_reg = (lazy_regularization and running_mb_counter % G_reg_interval == 0)
283
- run_D_reg = (lazy_regularization and running_mb_counter % D_reg_interval == 0)
284
- cur_nimg += sched.minibatch_size
285
- running_mb_counter += 1
286
-
287
- # Fast path without gradient accumulation.
288
- if len(rounds) == 1:
289
- tflib.run([G_train_op, data_fetch_op], feed_dict)
290
- if run_G_reg:
291
- tflib.run(G_reg_op, feed_dict)
292
- tflib.run([D_train_op, Gs_update_op], feed_dict)
293
- if run_D_reg:
294
- tflib.run(D_reg_op, feed_dict)
295
-
296
- # Slow path with gradient accumulation.
297
- else:
298
- for _round in rounds:
299
- tflib.run(G_train_op, feed_dict)
300
- if run_G_reg:
301
- for _round in rounds:
302
- tflib.run(G_reg_op, feed_dict)
303
- tflib.run(Gs_update_op, feed_dict)
304
- for _round in rounds:
305
- tflib.run(data_fetch_op, feed_dict)
306
- tflib.run(D_train_op, feed_dict)
307
- if run_D_reg:
308
- for _round in rounds:
309
- tflib.run(D_reg_op, feed_dict)
310
-
311
- # Perform maintenance tasks once per tick.
312
- done = (cur_nimg >= total_kimg * 1000)
313
- if cur_tick < 0 or cur_nimg >= tick_start_nimg + sched.tick_kimg * 1000 or done:
314
- cur_tick += 1
315
- tick_kimg = (cur_nimg - tick_start_nimg) / 1000.0
316
- tick_start_nimg = cur_nimg
317
- tick_time = dnnlib.RunContext.get().get_time_since_last_update()
318
- total_time = dnnlib.RunContext.get().get_time_since_start() + resume_time
319
-
320
- # Report progress.
321
- print('tick %-5d kimg %-8.1f lod %-5.2f minibatch %-4d time %-12s sec/tick %-7.1f sec/kimg %-7.2f maintenance %-6.1f gpumem %.1f' % (
322
- autosummary('Progress/tick', cur_tick),
323
- autosummary('Progress/kimg', cur_nimg / 1000.0),
324
- autosummary('Progress/lod', sched.lod),
325
- autosummary('Progress/minibatch', sched.minibatch_size),
326
- dnnlib.util.format_time(autosummary('Timing/total_sec', total_time)),
327
- autosummary('Timing/sec_per_tick', tick_time),
328
- autosummary('Timing/sec_per_kimg', tick_time / tick_kimg),
329
- autosummary('Timing/maintenance_sec', maintenance_time),
330
- autosummary('Resources/peak_gpu_mem_gb', peak_gpu_mem_op.eval() / 2**30)))
331
- autosummary('Timing/total_hours', total_time / (60.0 * 60.0))
332
- autosummary('Timing/total_days', total_time / (24.0 * 60.0 * 60.0))
333
-
334
- # Save snapshots.
335
- if image_snapshot_ticks is not None and (cur_tick % image_snapshot_ticks == 0 or done):
336
- grid_fakes = Gs.run(grid_latents, grid_labels, is_validation=True, minibatch_size=sched.minibatch_gpu)
337
- misc.save_image_grid(grid_fakes, dnnlib.make_run_dir_path('fakes%06d.png' % (cur_nimg // 1000)), drange=drange_net, grid_size=grid_size)
338
- if network_snapshot_ticks is not None and (cur_tick % network_snapshot_ticks == 0 or done):
339
- pkl = dnnlib.make_run_dir_path('network-snapshot-%06d.pkl' % (cur_nimg // 1000))
340
- misc.save_pkl((G, D, Gs), pkl)
341
- metrics.run(pkl, run_dir=dnnlib.make_run_dir_path(), data_dir=dnnlib.convert_path(data_dir), num_gpus=num_gpus, tf_config=tf_config)
342
-
343
- # Update summaries and RunContext.
344
- metrics.update_autosummaries()
345
- tflib.autosummary.save_summaries(summary_log, cur_nimg)
346
- dnnlib.RunContext.get().update('%.2f' % sched.lod, cur_epoch=cur_nimg // 1000, max_epoch=total_kimg)
347
- maintenance_time = dnnlib.RunContext.get().get_last_update_interval() - tick_time
348
-
349
- # Save final snapshot.
350
- misc.save_pkl((G, D, Gs), dnnlib.make_run_dir_path('network-final.pkl'))
351
-
352
- # All done.
353
- summary_log.close()
354
- training_set.close()
355
-
356
- #----------------------------------------------------------------------------