'training.networks.conditional.disc_filter_weight_mod_cutoff.D_stylegan2',
cutoff_layer=layer)
D.copy_vars_from(D_original)
for label_index, label_name in enumerate(labels_names):
# current_label = np.argmax(labels[0, offsets[label_name]: offsets[label_name] + num_classes[label_name]])
prediction, fmap_out = D.get_output_for(images_placeholder,
label_placeholder)
# gradient = tf.gradients(prediction[:, offsets[label_name] + current_label], [fmap_out])[0]
gradient = \
tf.gradients(prediction[:, offsets[label_name]:offsets[label_name] + num_classes[label_name]], [fmap_out])[
0]
prediction_out, fmap_out, gradient_out = tflib.run(
[prediction, fmap_out, gradient],
feed_dict={
label_placeholder: labels,
images_placeholder: images
})
a_c_k = np.mean(gradient_out[0], axis=(1, 2))
a_c_k_broadcast = np.tile(np.expand_dims(a_c_k, axis=-1),
[fmap_out.shape[2]])
a_c_k_broadcast = np.tile(np.expand_dims(a_c_k_broadcast, axis=-1),
[fmap_out.shape[3]])
result = relu(a_c_k_broadcast * fmap_out[0])
result_negative = relu_inverse(a_c_k_broadcast * fmap_out[0])
result = np.sum(result, axis=0)
result_negative = np.sum(result_negative, axis=0)
mask = result + result_negative
# max_mask = np.max(mask)
# min_mask = np.min(mask)
def images_uint8(self):
return tflib.run(self._images_uint8_expr, {self._dlatent_noise_in: 0})
def main():
"""Main function."""
args = parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
assert os.path.exists(args.image_list)
image_list_name = os.path.splitext(os.path.basename(args.image_list))[0]
output_dir = args.output_dir or f'results/ghfeat/{image_list_name}'
logger = setup_logger(output_dir, 'extract_feature.log',
'inversion_logger')
logger.info(f'Loading model.')
tflib.init_tf({'rnd.np_random_seed': 1000})
with open(args.model_path, 'rb') as f:
E, _, _, Gs = pickle.load(f)
# Get input size.
image_size = E.input_shape[2]
assert image_size == E.input_shape[3]
G_args = EasyDict(func_name='training.networks_stylegan.G_synthesis')
G_style_mod = tflib.Network('G_StyleMod',
resolution=image_size,
label_size=0,
**G_args)
Gs_vars_pairs = {
name: tflib.run(val)
for name, val in Gs.components.synthesis.vars.items()
}
for g_name, g_val in G_style_mod.vars.items():
tflib.set_vars({g_val: Gs_vars_pairs[g_name]})
# Build graph.
logger.info(f'Building graph.')
sess = tf.get_default_session()
input_shape = E.input_shape
input_shape[0] = args.batch_size
x = tf.placeholder(tf.float32, shape=input_shape, name='real_image')
ghfeat = E.get_output_for(x, is_training=False)
x_rec = G_style_mod.get_output_for(ghfeat, randomize_noise=False)
# Load image list.
logger.info(f'Loading image list.')
image_list = []
with open(args.image_list, 'r') as f:
for line in f:
image_list.append(line.strip())
# Extract GH-Feat from images.
logger.info(f'Start feature extraction.')
headers = ['Name', 'Original Image', 'Encoder Output']
viz_size = None if args.viz_size == 0 else args.viz_size
visualizer = HtmlPageVisualizer(num_rows=len(image_list),
num_cols=len(headers),
viz_size=viz_size)
visualizer.set_headers(headers)
images = np.zeros(input_shape, np.uint8)
names = ['' for _ in range(args.batch_size)]
features = []
for img_idx in tqdm(range(0, len(image_list), args.batch_size),
leave=False):
# Load inputs.
batch = image_list[img_idx:img_idx + args.batch_size]
for i, image_path in enumerate(batch):
image = resize_image(load_image(image_path),
(image_size, image_size))
images[i] = np.transpose(image, [2, 0, 1])
names[i] = os.path.splitext(os.path.basename(image_path))[0]
inputs = images.astype(np.float32) / 255 * 2.0 - 1.0
# Run encoder.
outputs = sess.run([ghfeat, x_rec], {x: inputs})
features.append(outputs[0][0:len(batch)])
outputs[1] = adjust_pixel_range(outputs[1])
for i, _ in enumerate(batch):
image = np.transpose(images[i], [1, 2, 0])
save_image(f'{output_dir}/{names[i]}_ori.png', image)
save_image(f'{output_dir}/{names[i]}_enc.png', outputs[1][i])
visualizer.set_cell(i + img_idx, 0, text=names[i])
visualizer.set_cell(i + img_idx, 1, image=image)
visualizer.set_cell(i + img_idx, 2, image=outputs[1][i])
# Save results.
os.system(f'cp {args.image_list} {output_dir}/image_list.txt')
np.save(f'{output_dir}/ghfeat.npy', np.concatenate(features, axis=0))
visualizer.save(f'{output_dir}/reconstruction.html')
label[:, :108], rotations[:, (rotation + 1) % 8],
label[:, 108 + 8:]
],
axis=-1)
label_interpolate = label_left * (
1 - magnitude) + label_right * magnitude
# print(label_interpolate[0, 108:116])
# change background to city for kitti dataset
if only_city_background:
# city_label = np.array([[0, 1, 0, 0, 0, 0]]) # city
city_label = np.array([[0, 0, 0, 1, 0, 0]])
label_interpolate = np.concatenate(
[label_interpolate[:, :-6], city_label], axis=-1)
out = tflib.run(synthesise,
feed_dict={label_placeholder: label_interpolate})
image = convert_to_image(out[0])
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
if target_angle >= 180 and flip_images:
image = cv2.flip(image, 1)
predicted_angle = detector.predict_angle(image, show_crop=False)
if predicted_angle:
if target_angle >= 180 and flip_images:
predicted_angle = 360 - predicted_angle
print(target_angle, predicted_angle)
predicted_angle = detector.correct_angles(
predicted_angle,
target_angle,
correct_threshold=180,
correct_angles=[360])
def noises(self):
return tflib.run(self._noise_vars)
def training_loop(
G_args = {}, # Options for generator network.
D_args = {}, # Options for discriminator network.
G_opt_args = {}, # Options for generator optimizer.
D_opt_args = {}, # Options for discriminator optimizer.
G_loss_args = {}, # Options for generator loss.
D_loss_args = {}, # Options for discriminator loss.
dataset_args = {}, # Options for dataset.load_dataset().
sched_args = {}, # Options for train.TrainingSchedule.
grid_args = {}, # Options for train.setup_snapshot_image_grid().
metric_arg_list = [], # Options for MetricGroup.
tf_config = {}, # Options for tflib.init_tf().
data_dir = None, # Directory to load datasets from.
G_smoothing_kimg = 10.0, # Half-life of the running average of generator weights.
minibatch_repeats = 4, # Number of minibatches to run before adjusting training parameters.
lazy_regularization = True, # Perform regularization as a separate training step?
G_reg_interval = 4, # How often the perform regularization for G? Ignored if lazy_regularization=False.
D_reg_interval = 16, # How often the perform regularization for D? Ignored if lazy_regularization=False.
reset_opt_for_new_lod = True, # Reset optimizer internal state (e.g. Adam moments) when new layers are introduced?
total_kimg = 25000, # Total length of the training, measured in thousands of real images.
mirror_augment = False, # Enable mirror augment?
drange_net = [0,1], # Dynamic range used when feeding image data to the networks.
image_snapshot_ticks = 50, # How often to save image snapshots? None = only save 'reals.png' and 'fakes-init.png'.
network_snapshot_ticks = 50, # How often to save network snapshots? None = only save 'networks-final.pkl'.
save_tf_graph = False, # Include full TensorFlow computation graph in the tfevents file?
save_weight_histograms = False, # Include weight histograms in the tfevents file?
resume_pkl = None, # Network pickle to resume training from, None = train from scratch.
resume_kimg = 0.0, # Assumed training progress at the beginning. Affects reporting and training schedule.
resume_time = 0.0, # Assumed wallclock time at the beginning. Affects reporting.
resume_with_new_nets = False): # Construct new networks according to G_args and D_args before resuming training?
# Initialize dnnlib and TensorFlow.
tflib.init_tf(tf_config)
num_gpus = dnnlib.submit_config.num_gpus
# Load training set.
training_set = dataset.load_dataset(data_dir=dnnlib.convert_path(data_dir), verbose=True, **dataset_args)
grid_size, grid_reals, grid_labels = misc.setup_snapshot_image_grid(training_set, **grid_args)
misc.save_image_grid(grid_reals, dnnlib.make_run_dir_path('reals.png'), drange=training_set.dynamic_range, grid_size=grid_size)
# Construct or load networks.
training_set.configure(minibatch_size=sched_args.batch_size)
with tf.device('/gpu:0'):
if resume_pkl is None or resume_with_new_nets:
print('Constructing networks...')
G = tflib.Network('G', num_channels=training_set.shape[0], resolution=training_set.shape[1], label_size=training_set.label_size, **G_args)
D = tflib.Network('D', num_channels=training_set.shape[0], resolution=training_set.shape[1], label_size=training_set.label_size, **D_args)
Gs = G.clone('Gs')
start = 0
if resume_pkl is not None:
print('Loading networks from "%s"...' % resume_pkl)
rG, rD, rGs = misc.load_pkl(resume_pkl)
if resume_with_new_nets: G.copy_vars_from(rG); D.copy_vars_from(rD); Gs.copy_vars_from(rGs)
else: G = rG; D = rD; Gs = rGs
start = int(resume_pkl.split('-')[-1].split('.')[0]) // sched_args.batch_size
# Print layers and generate initial image snapshot.
G.print_layers(); D.print_layers()
grid_latents = np.random.randn(np.prod(grid_size), *G.input_shape[1:])
grid_fakes = G.run(grid_latents, grid_labels, is_validation=True, minibatch_size=sched_args.batch_size)
misc.save_image_grid(grid_fakes, dnnlib.make_run_dir_path('fakes_init.png'), drange=drange_net, grid_size=grid_size)
global_step = tf.Variable(start, trainable=False, name='learning_rate_step')
learning_rate = tf.train.exponential_decay(sched_args.lr, global_step, sched_args.decay_step,
sched_args.decay_rate, staircase=sched_args.stair)
add_global = global_step.assign_add(1)
D_opt = tflib.Optimizer(name='TrainD', learning_rate=learning_rate, **D_opt_args)
G_opt = tflib.Optimizer(name='TrainG', learning_rate=learning_rate, **G_opt_args)
for gpu in range(num_gpus):
print('build graph on gpu %s' % str(gpu))
with tf.name_scope('GPU%d' % gpu), tf.device('/gpu:%d' % gpu):
# Create GPU-specific shadow copies of G and D.
G_gpu = G if gpu == 0 else G.clone(G.name + '_shadow')
D_gpu = D if gpu == 0 else D.clone(D.name + '_shadow')
with tf.name_scope('DataFetch'):
reals_read, labels_read = training_set.get_minibatch_tf()
reals_read, labels_read = process_reals(reals_read, labels_read, mirror_augment,
training_set.dynamic_range, drange_net)
with tf.name_scope('Loss'), tf.control_dependencies(None):
loss, reg = dnnlib.util.call_func_by_name(G=G_gpu, D=D_gpu, opt=D_opt, training_set=training_set,
minibatch_size=sched_args.batch_size, reals=reals_read,
labels=labels_read, **D_loss_args)
with tf.control_dependencies([add_global]):
G_opt.register_gradients(loss, G_gpu.trainables)
D_opt.register_gradients(loss, D_gpu.trainables)
G_train_op = G_opt.apply_updates()
D_train_op = D_opt.apply_updates()
# Finalize graph.
with tf.device('/gpu:0'):
try:
peak_gpu_mem_op = tf.contrib.memory_stats.MaxBytesInUse()
except tf.errors.NotFoundError:
peak_gpu_mem_op = tf.constant(0)
tflib.init_uninitialized_vars()
print('Initializing logs...')
summary_log = tf.summary.FileWriter(dnnlib.make_run_dir_path())
if save_tf_graph:
summary_log.add_graph(tf.get_default_graph())
if save_weight_histograms:
G.setup_weight_histograms(); D.setup_weight_histograms()
metrics = metric_base.MetricGroup(metric_arg_list)
print('Training for %d kimg...\n' % total_kimg)
dnnlib.RunContext.get().update('', cur_epoch=resume_kimg, max_epoch=total_kimg)
maintenance_time = dnnlib.RunContext.get().get_last_update_interval()
cur_nimg = int(resume_kimg * 1000)
cur_tick = -1
tick_start_nimg = cur_nimg
prev_lod = -1.0
running_mb_counter = 0
while cur_nimg < total_kimg * 1000:
if dnnlib.RunContext.get().should_stop(): break
loss_, _, _, lr_ = tflib.run([loss, G_train_op, D_train_op, learning_rate])
cur_nimg += sched_args.batch_size * num_gpus
done = (cur_nimg >= total_kimg * 1000)
if cur_tick < 0 or cur_nimg >= tick_start_nimg + sched_args.tick_kimg * 1000 or done:
cur_tick += 1
tick_kimg = (cur_nimg - tick_start_nimg) / 1000.0
tick_start_nimg = cur_nimg
tick_time = dnnlib.RunContext.get().get_time_since_last_update()
total_time = dnnlib.RunContext.get().get_time_since_start() + resume_time
# Report progress.
print(
'tick %-5d kimg %-8.1f minibatch %-4d time %-12s sec/tick %-7.1f '
'sec/kimg %-7.2f maintenance %-6.1f gpumem %.1f loss %-8.1f lr %-2.5f' % (
autosummary('Progress/tick', cur_tick),
autosummary('Progress/kimg', cur_nimg / 1000.0),
autosummary('Progress/minibatch', sched_args.batch_size),
dnnlib.util.format_time(autosummary('Timing/total_sec', total_time)),
autosummary('Timing/sec_per_tick', tick_time),
autosummary('Timing/sec_per_kimg', tick_time / tick_kimg),
autosummary('Timing/maintenance_sec', maintenance_time),
autosummary('Resources/peak_gpu_mem_gb', peak_gpu_mem_op.eval() / 2 ** 30),
autosummary('loss', loss_),
autosummary('lr', lr_)))
autosummary('Timing/total_hours', total_time / (60.0 * 60.0))
autosummary('Timing/total_days', total_time / (24.0 * 60.0 * 60.0))
# Save snapshots.
if image_snapshot_ticks is not None and (cur_tick % image_snapshot_ticks == 0 or done):
grid_fakes = G.run(grid_latents, grid_labels, is_validation=True, minibatch_size=sched_args.batch_size)
misc.save_image_grid(grid_fakes, dnnlib.make_run_dir_path('fakes%06d.png' % (cur_nimg // 1000)),
drange=drange_net, grid_size=grid_size)
if network_snapshot_ticks is not None and (cur_tick % network_snapshot_ticks == 0 or done):
pkl = dnnlib.make_run_dir_path('network-snapshot-%06d.pkl' % (cur_nimg // 1000))
misc.save_pkl((G, D, Gs), pkl)
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)
# Update summaries and RunContext.
metrics.update_autosummaries()
tflib.autosummary.save_summaries(summary_log, cur_nimg)
dnnlib.RunContext.get().update('%.2f' % 0.0, cur_epoch=cur_nimg // 1000, max_epoch=total_kimg)
maintenance_time = dnnlib.RunContext.get().get_last_update_interval() - tick_time
# Save final snapshot.
misc.save_pkl((G, D, Gs), dnnlib.make_run_dir_path('network-final.pkl'))
# All done.
summary_log.close()
training_set.close()
def get_minibatch_np(self, minibatch_size, lod=0): # => images, labels
self.configure(minibatch_size, lod)
if self._tf_minibatch_np is None:
self._tf_minibatch_np = self.get_minibatch_tf()
return tflib.run(self._tf_minibatch_np)
condition = tf.equal(labels[:, 108], 0.7071)
random_vector = tf.random_uniform([minibatch_size]) < 0.5
remove_condition = tf.logical_and(condition, random_vector)
labels = tf.where(remove_condition, removed_labels, labels)
random_angle = tf.random_uniform([minibatch_size]) * 2 * np.pi
interpolation_rotation_cos = tf.expand_dims(tf.cos(random_angle), axis=-1)
interpolation_rotation_sin = tf.expand_dims(tf.sin(random_angle), axis=-1)
interpolated_labels = tf.concat([
labels[:, :rotation_offset],
interpolation_rotation_cos,
interpolation_rotation_sin,
labels[:, rotation_offset + 2:]
], axis=1)
interpolation_mask = tf.random_uniform([minibatch_size]) < (1 - interpolation_prob)
labels_mixed = tf.where(interpolation_mask, labels, interpolated_labels)
fig, axes = plt.subplots(subplot_kw=dict(projection='polar'))
for i in range(10):
labels_out = tflib.run(labels_mixed)
rotation_cos = labels_out[:, rotation_offset]
rotation_sin = labels_out[:, rotation_offset + 1]
distance = np.linalg.norm([rotation_cos, rotation_sin], ord=2)
angle = np.arctan2(rotation_sin, rotation_cos)
axes.scatter(angle, distance, c='red')
plt.show()
dtype=tf.float32)
indices = tf.cast(tf.floor(
tf.random_uniform(shape=[minibatch_size], minval=0, maxval=8)),
dtype=tf.int32)
random_rotation = tf.gather(all_rotations, indices)
labels = tf.concat([
labels[:, :rotation_offset], random_rotation, labels[:,
rotation_offset + 2:]
],
axis=1)
random_angle = tf.random_uniform([tf.shape(x)[0]]) * 2 * np.pi
interpolation_rotation_cos = tf.expand_dims(tf.cos(random_angle), axis=-1)
interpolation_rotation_sin = tf.expand_dims(tf.sin(random_angle), axis=-1)
interpolated_labels = tf.concat([
labels[:, :rotation_offset], interpolation_rotation_cos,
interpolation_rotation_sin, labels[:, rotation_offset + 2:]
],
axis=1)
interpolation_mask = tf.random_uniform([tf.shape(x)[0]]) < 0.5
labels = tf.where(interpolation_mask, labels, interpolated_labels)
labels_out, random_rotation, indices = tflib.run(
[labels, random_rotation, indices])
print(labels_out[:, 108:110])
print(np.arctan2(labels_out[:, 108], labels_out[:, 109]) / (2 * np.pi) * 360)
print(np.linalg.norm(labels_out[:, 108:110], axis=-1, ord=2))
def get_minibatch_np(self, minibatch_size, lod=0): # => images, labels
self.configure(minibatch_size, lod)
with tf.name_scope('Dataset'):
if self._tf_minibatch_np is None:
self._tf_minibatch_np = self.get_minibatch_tf()
return tflib.run(self._tf_minibatch_np)
def _evaluate(self, Gs, E, Inv, num_gpus):
minibatch_size = num_gpus * self.minibatch_per_gpu
resolution = Gs.components.synthesis.output_shape[2]
placeholder_portraits = tf.placeholder(
tf.float32, [self.minibatch_per_gpu, 3, resolution, resolution],
name='placeholder_portraits')
placeholder_landmarks = tf.placeholder(
tf.float32, [self.minibatch_per_gpu, 3, resolution, resolution],
name='placeholder_landmarks')
placeholder_keypoints = tf.placeholder(tf.float32,
[self.minibatch_per_gpu, 136],
name='placeholder_landmarks')
fake_X_val = self.run_image_manipulation(E, Gs, Inv,
placeholder_portraits,
placeholder_landmarks,
placeholder_keypoints,
num_gpus)
csim_sum = 0.0
for idx, data in enumerate(
self._iterate_reals(minibatch_size=minibatch_size)):
image_data = data[0]
batch_portraits = image_data[:, 0, :, :, :]
batch_landmarks = np.roll(image_data[:, 1, :, :, :],
shift=1,
axis=0)
keypoints = np.roll(data[1], shift=1, axis=0)
batch_portraits = misc.adjust_dynamic_range(
batch_portraits.astype(np.float32), [0, 255], [-1., 1.])
batch_landmarks = misc.adjust_dynamic_range(
batch_landmarks.astype(np.float32), [0, 255], [-1., 1.])
begin = idx * minibatch_size
end = min(begin + minibatch_size, self.num_images)
samples_manipulated = tflib.run(fake_X_val,
feed_dict={
placeholder_portraits:
batch_portraits,
placeholder_landmarks:
batch_landmarks,
placeholder_keypoints:
keypoints
})
samples_manipulated = np.transpose(samples_manipulated,
[0, 2, 3, 1])
samples_manipulated = np.pad(samples_manipulated,
((0, 0), (11, 11), (11, 11), (0, 0)),
mode='constant')
batch_portraits = np.transpose(batch_portraits, [0, 2, 3, 1])
batch_portraits = np.pad(batch_portraits,
((0, 0), (11, 11), (11, 11), (0, 0)),
mode='constant')
embeddings_real = self.get_facenet_embeddings(batch_portraits)
embeddings_fake = self.get_facenet_embeddings(samples_manipulated)
for i in range(minibatch_size):
csim_sum += 1 - scipy.spatial.distance.cosine(
embeddings_real[i], embeddings_fake[i])
if end == self.num_images:
break
avg_csim = csim_sum / self.num_images
self._report_result(np.real(avg_csim))
def get_clean_images(self):
return tflib.run(self._images_clean, {self._noise_in: 0})
def step(self):
assert self._cur_step is not None
if self._cur_step >= self.num_steps:
return
if self._cur_step == 0:
self._info('Running...')
# Hyperparameters.
t = self._cur_step / self.num_steps
noise_strength = self._dlatent_std * self.initial_noise_factor * max(
0.0, 1.0 - t / self.noise_ramp_length)**2
lr_ramp = min(1.0, (1.0 - t) / self.lr_rampdown_length)
lr_ramp = 0.5 - 0.5 * np.cos(lr_ramp * np.pi)
lr_ramp = lr_ramp * min(1.0, t / self.lr_rampup_length)
learning_rate = self.initial_learning_rate * lr_ramp
#noise_reg_curr = 0
#l2_pixelwise_reg = 0
l2_pixelwise_reg = self.l2_pixelwise_reg
if t < 1: # change to 0.5 when noise optim starts working
noise_reg_curr = 0
#l2_pixelwise_reg = 0
reg_dlatent_fixed = 0
else:
noise_reg_curr = 1
#l2_pixelwise_reg = self.l2_pixelwise_reg
# set fixed dlatent, and copy the value into the fixed dlatent variable
reg_dlatent_fixed = 1
#tflib.run(tf.assign(self._dlatents_expr_fixed, self._dlatents_expr_moving))
#if t >= 0.2 and t < 0.21:
# self._dlatents_expr.require_grad = False
# Train.
feed_dict = {
self._noise_in: noise_strength,
self._lrate_in: learning_rate,
self._l2_pixelwise_reg: l2_pixelwise_reg,
self.noise_reg: noise_reg_curr,
self.reg_dlatent_fixed: reg_dlatent_fixed
}
_, dist_value, _loss, _pixelwise_loss, _w_loss, _cosine_loss, _perceptual_loss = tflib.run(
[
self._opt_step, self._dist, self._loss, self._pixelwise_loss,
self._w_loss, self._cosine_loss, self._perceptual_loss
], feed_dict)
#tflib.run(self._noise_normalize_op)
# Print status.
self._cur_step += 1
if self._cur_step == self.num_steps or self._cur_step % 50 == 0:
print('_loss:', _loss, flush=True)
print('_perceptual:', _perceptual_loss, flush=True)
print('_pixelwise:', _pixelwise_loss, flush=True)
print('_w_loss:', _w_loss, flush=True)
print('_cosine:', _cosine_loss, flush=True)
self._info('%-8d%-12g%-12g' % (self._cur_step, dist_value, _loss))
if self._cur_step == self.num_steps:
self._info('Done.')
def start(self, target_images):
assert self._Gs is not None
# Prepare target images.
self._info('Preparing target images...')
target_images = np.asarray(target_images, dtype='float32')
target_images = (target_images + 1) * (255 / 2)
sh = target_images.shape
assert sh[0] == self._minibatch_size
if vggDownscale:
#if sh[2] > self._target_images_var.shape[2]:
print('self._target_images_var', self._target_images_var.shape)
print('self._target_images_down_var',
self._target_images_down_var.shape)
if self._target_images_down_var.shape[
2] < self._target_images_var.shape[2]:
print('checkpoint vgg downscale')
factor = sh[2] // self._target_images_down_var.shape[2]
target_images_down = np.reshape(target_images, [
-1, sh[1], sh[2] // factor, factor, sh[3] // factor, factor
]).mean((3, 5))
else:
target_images_down = target_images
if sh[1] == 1:
# if mono images, make them 3-channel, so you can use VGG to compute LPIPS distance. VGG needs 3-channel images as input
target_images_down = np.repeat(target_images_down, 3, axis=1)
target_images = np.repeat(target_images, 3, axis=1)
print('target_images', target_images.shape)
print('target_images_down', target_images_down.shape)
print('target_images_down_var', self._target_images_down_var.shape)
# Initialize optimization state.
self._info('Initializing optimization state...')
tflib.set_vars({
self._target_images_var:
target_images,
self._target_images_down_var:
target_images_down,
self._dlatents_expr_moving:
np.tile(self._dlatent_avg,
[self._minibatch_size, self.numLatentLayers, 1]),
self._dlatents_expr_fixed:
np.zeros(self.dlatent_shape),
self._l2_pixelwise_reg:
0.0,
self.noise_reg:
0.0,
self.reg_dlatent_fixed:
0
})
self.forward.initVars()
print('start() target_images shape', target_images.shape)
#print('mask True', np.sum(self.forward.mask))
#print('mask False', np.sum(~self.forward.mask))
#self.forward.calcMaskFromImg(target_images)
#asdad
tflib.run(self._noise_init_op)
#tflib.run(tf.group(tf.assign(self._l2_pixelwise_reg, 0)))
self._opt.reset_optimizer_state()
self._cur_step = 0
def get_images(self):
return tflib.run(self._images_expr, {self._noise_in: 0})
def _evaluate(self, classifier, Gs_kwargs, num_gpus):
self._set_dataset_obj(
dataset.TFRecordDataset(tfrecord_dir='datasets/classifier_dataset_15_test'))
iterations = self.num_images // self.minibatch_per_gpu
distance_list = []
images_placeholder = tf.placeholder(shape=classifier.input_shapes[0], dtype=tf.float32)
num_correct_model = 0
num_correct_color = 0
num_correct_manufacturer = 0
num_correct_body = 0
num_correct_rotation = 0
num_correct_ratio = 0
num_correct_background = 0
num_label_model = 0
num_label_color = 0
num_label_manufacturer = 0
num_label_body = 0
num_label_rotation = 0
num_label_ratio = 0
num_label_background = 0
for _ in range(iterations):
reals = next(self._iterate_reals(self.minibatch_per_gpu))
images, labels = reals
images = misc.adjust_dynamic_range(images, [0, 255], [-1, 1])
model_pred, color_pred, manufacturer_pred, body_pred, rotation_pred, ratio_pred, background_pred = classifier.get_output_for(
images_placeholder)
i = 0
offsets = [1, 67, 12, 18, 10, 8, 5, 6]
current_offset = offsets[i]
next_offset = current_offset + offsets[i + 1]
current_label = labels[:, current_offset:next_offset]
num_model_pred = tf.reduce_sum(tf.one_hot(indices=tf.argmax(model_pred, axis=-1), depth=offsets[i + 1]) * current_label)
num_model_label = tf.reduce_sum(current_label)
i += 1
current_offset = next_offset
next_offset = current_offset + offsets[i + 1]
current_label = labels[:, current_offset:next_offset]
num_color_pred = tf.reduce_sum(tf.one_hot(indices=tf.argmax(color_pred, axis=-1), depth=offsets[i + 1]) * current_label)
num_color_label = tf.reduce_sum(current_label)
i += 1
current_offset = next_offset
next_offset = current_offset + offsets[i + 1]
current_label = labels[:, current_offset:next_offset]
num_manufacturer_pred = tf.reduce_sum(tf.one_hot(indices=tf.argmax(manufacturer_pred, axis=-1), depth=offsets[i + 1]) * current_label)
num_manufacturer_label = tf.reduce_sum(current_label)
i += 1
current_offset = next_offset
next_offset = current_offset + offsets[i + 1]
current_label = labels[:, current_offset:next_offset]
num_body_pred = tf.reduce_sum(tf.one_hot(indices=tf.argmax(body_pred, axis=-1), depth=offsets[i + 1]) * current_label)
num_body_label = tf.reduce_sum(current_label)
i += 1
current_offset = next_offset
next_offset = current_offset + offsets[i + 1]
current_label = labels[:, current_offset:next_offset]
num_rotation_pred = tf.reduce_sum(tf.one_hot(indices=tf.argmax(rotation_pred, axis=-1), depth=offsets[i + 1]) * current_label)
num_rotation_label = tf.reduce_sum(current_label)
i += 1
current_offset = next_offset
next_offset = current_offset + offsets[i + 1]
current_label = labels[:, current_offset:next_offset]
num_ratio_pred = tf.reduce_sum(tf.one_hot(indices=tf.argmax(ratio_pred, axis=-1), depth=offsets[i + 1]) * current_label)
num_ratio_label = tf.reduce_sum(current_label)
i += 1
current_offset = next_offset
next_offset = current_offset + offsets[i + 1]
current_label = labels[:, current_offset:next_offset]
num_background_pred = tf.reduce_sum(tf.one_hot(indices=tf.argmax(background_pred, axis=-1), depth=offsets[i + 1]) * current_label)
num_background_label = tf.reduce_sum(current_label)
output = tflib.run([
num_model_pred,
num_color_pred,
num_manufacturer_pred,
num_body_pred,
num_rotation_pred,
num_ratio_pred,
num_background_pred,
num_model_label,
num_color_label,
num_manufacturer_label,
num_body_label,
num_rotation_label,
num_ratio_label,
num_background_label
], feed_dict={images_placeholder: images})
num_correct_model += output[0]
num_correct_color += output[1]
num_correct_manufacturer += output[2]
num_correct_body += output[3]
num_correct_rotation += output[4]
num_correct_ratio += output[5]
num_correct_background += output[6]
num_label_model += output[7]
num_label_color += output[8]
num_label_manufacturer += output[9]
num_label_body += output[10]
num_label_rotation += output[11]
num_label_ratio += output[12]
num_label_background += output[13]
model_acc = num_correct_model / num_label_model
print('Model Accuracy: {}'.format(model_acc))
color_acc = num_correct_color / num_label_color
print('Color Accuracy: {}'.format(color_acc))
manufacturer_acc = num_correct_manufacturer / num_label_manufacturer
print('Manufacturer Accuracy: {}'.format(manufacturer_acc))
body_acc = num_correct_body / num_label_body
print('Body Accuracy: {}'.format(body_acc))
rotation_acc = num_correct_rotation / num_label_rotation
print('Rotation Accuracy: {}'.format(rotation_acc))
ratio_acc = num_correct_ratio / num_label_ratio
print('Ratio Accuracy: {}'.format(ratio_acc))
background_acc = num_correct_background / num_label_background
print('Background Accuracy: {}'.format(background_acc))
self._report_result(np.mean([
model_acc,
color_acc,
manufacturer_acc,
body_acc,
rotation_acc,
ratio_acc,
background_acc
]))
def get_lpips(self):
return tflib.run(self._lpips_loss, {self._noise_in: 0})
for angle in range(0, 360, 12):
angles.append(angle)
angle_pi = (angle / 360) * 2 * np.pi
# change rotation
cos = np.expand_dims(np.expand_dims(np.cos(angle_pi), axis=0), axis=0)
sin = np.expand_dims(np.expand_dims(np.sin(angle_pi), axis=0), axis=0)
label_adjust = np.concatenate(
[label[:, :108], cos, sin, label[:, 108 + 2:]], axis=-1)
# change background to city for kitti dataset
city_label = np.array([[0, 1, 0, 0, 0, 0]])
label_adjust = np.concatenate([label_adjust[:, :-6], city_label],
axis=-1)
out = tflib.run(synthesise,
feed_dict={label_placeholder: label_adjust})
image = convert_to_image(out[0])
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
if angle >= 180:
image = cv2.flip(image, 1)
img = np.copy(image)
yolo_img = np.copy(image)
detections = yolo.detect(yolo_img)
detections = filter(
lambda detection: detection.detected_class == 'car', detections)
max_size = 0
keep_detection = None
# only keep the largest car detection
def training_loop(
# Configurations
cG={},
cD={}, # Generator and Discriminator command-line arguments
dataset_args={}, # dataset.load_dataset() options
sched_args={}, # train.TrainingSchedule options
vis_args={}, # vis.eval options
grid_args={}, # train.setup_snapshot_img_grid() options
metric_arg_list=[], # MetricGroup Options
tf_config={}, # tflib.init_tf() options
eval=False, # Evaluation mode
train=False, # Training mode
# Data
data_dir=None, # Directory to load datasets from
total_kimg=25000, # Total length of the training, measured in thousands of real images
mirror_augment=False, # Enable mirror augmentation?
drange_net=[
-1, 1
], # Dynamic range used when feeding image data to the networks
ratio=1.0, # Image height/width ratio in the dataset
# Optimization
minibatch_repeats=4, # Number of minibatches to run before adjusting training parameters
lazy_regularization=True, # Perform regularization as a separate training step?
smoothing_kimg=10.0, # Half-life of the running average of generator weights
clip=None, # Clip gradients threshold
# Resumption
resume_pkl=None, # Network pickle to resume training from, None = train from scratch.
resume_kimg=0.0, # Assumed training progress at the beginning
# Affects reporting and training schedule
resume_time=0.0, # Assumed wallclock time at the beginning, affects reporting
recompile=False, # Recompile network from source code (otherwise loads from snapshot)
# Logging
summarize=True, # Create TensorBoard summaries
save_tf_graph=False, # Include full TensorFlow computation graph in the tfevents file?
save_weight_histograms=False, # Include weight histograms in the tfevents file?
img_snapshot_ticks=3, # How often to save image snapshots? None = disable
network_snapshot_ticks=3, # How often to save network snapshots? None = only save networks-final.pkl
last_snapshots=10, # Maximal number of prior snapshots to save
eval_images_num=50000, # Sample size for the metrics
printname="", # Experiment name for logging
# Architecture
merge=False): # Generate several images and then merge them
# Initialize dnnlib and TensorFlow
tflib.init_tf(tf_config)
num_gpus = dnnlib.submit_config.num_gpus
cG.name, cD.name = "g", "d"
# Load dataset, configure training scheduler and metrics object
dataset = data.load_dataset(data_dir=dnnlib.convert_path(data_dir),
verbose=True,
**dataset_args)
sched = training_schedule(sched_args,
cur_nimg=total_kimg * 1000,
dataset=dataset)
metrics = metric_base.MetricGroup(metric_arg_list)
# Construct or load networks
with tf.device("/gpu:0"):
no_op = tf.no_op()
G, D, Gs = None, None, None
if resume_pkl is None or recompile:
misc.log("Constructing networks...", "white")
G = tflib.Network("G",
num_channels=dataset.shape[0],
resolution=dataset.shape[1],
label_size=dataset.label_size,
**cG.args)
D = tflib.Network("D",
num_channels=dataset.shape[0],
resolution=dataset.shape[1],
label_size=dataset.label_size,
**cD.args)
Gs = G.clone("Gs")
if resume_pkl is not None:
G, D, Gs = load_nets(resume_pkl, G, D, Gs, recompile)
G.print_layers()
D.print_layers()
# Train/Evaluate/Visualize
# Labels are optional but not essential
grid_size, grid_reals, grid_labels = misc.setup_snapshot_img_grid(
dataset, **grid_args)
misc.save_img_grid(grid_reals,
dnnlib.make_run_dir_path("reals.png"),
drange=dataset.dynamic_range,
grid_size=grid_size)
grid_latents = np.random.randn(np.prod(grid_size), *G.input_shape[1:])
if eval:
# Save a snapshot of the current network to evaluate
pkl = dnnlib.make_run_dir_path("network-eval-snapshot-%06d.pkl" %
resume_kimg)
misc.save_pkl((G, D, Gs), pkl, remove=False)
# Quantitative evaluation
misc.log("Run evaluation...")
metric = metrics.run(pkl,
num_imgs=eval_images_num,
run_dir=dnnlib.make_run_dir_path(),
data_dir=dnnlib.convert_path(data_dir),
num_gpus=num_gpus,
ratio=ratio,
tf_config=tf_config,
eval_mod=True,
mirror_augment=mirror_augment)
# Qualitative evaluation
misc.log("Produce visualizations...")
visualize.eval(Gs,
dataset,
batch_size=sched.minibatch_gpu,
drange_net=drange_net,
ratio=ratio,
**vis_args)
if not train:
dataset.close()
exit()
# Setup training inputs
misc.log("Building TensorFlow graph...", "white")
with tf.name_scope("Inputs"), tf.device("/cpu:0"):
lrate_in_g = tf.placeholder(tf.float32, name="lrate_in_g", shape=[])
lrate_in_d = tf.placeholder(tf.float32, name="lrate_in_d", shape=[])
step = tf.placeholder(tf.int32, name="step", shape=[])
minibatch_size_in = tf.placeholder(tf.int32,
name="minibatch_size_in",
shape=[])
minibatch_gpu_in = tf.placeholder(tf.int32,
name="minibatch_gpu_in",
shape=[])
minibatch_multiplier = minibatch_size_in // (minibatch_gpu_in *
num_gpus)
beta = 0.5**tf.div(tf.cast(minibatch_size_in,
tf.float32), smoothing_kimg *
1000.0) if smoothing_kimg > 0.0 else 0.0
# Set optimizers
for cN, lr in [(cG, lrate_in_g), (cD, lrate_in_d)]:
set_optimizer(cN, lr, minibatch_multiplier, lazy_regularization, clip)
# Build training graph for each GPU
data_fetch_ops = []
for gpu in range(num_gpus):
with tf.name_scope("GPU%d" % gpu), tf.device("/gpu:%d" % gpu):
# Create GPU-specific shadow copies of G and D
for cN, N in [(cG, G), (cD, D)]:
cN.gpu = N if gpu == 0 else N.clone(N.name + "_shadow")
Gs_gpu = Gs if gpu == 0 else Gs.clone(Gs.name + "_shadow")
# Fetch training data via temporary variables
with tf.name_scope("DataFetch"):
reals, labels = dataset.get_minibatch_tf()
reals = process_reals(reals, dataset.dynamic_range, drange_net,
mirror_augment)
reals, reals_fetch = read_data(
reals, "reals", [sched.minibatch_gpu] + dataset.shape,
minibatch_gpu_in)
labels, labels_fetch = read_data(
labels, "labels",
[sched.minibatch_gpu, dataset.label_size],
minibatch_gpu_in)
data_fetch_ops += [reals_fetch, labels_fetch]
# Evaluate loss functions
with tf.name_scope("G_loss"):
cG.loss, cG.reg = dnnlib.util.call_func_by_name(
G=cG.gpu,
D=cD.gpu,
dataset=dataset,
reals=reals,
minibatch_size=minibatch_gpu_in,
**cG.loss_args)
with tf.name_scope("D_loss"):
cD.loss, cD.reg = dnnlib.util.call_func_by_name(
G=cG.gpu,
D=cD.gpu,
dataset=dataset,
reals=reals,
labels=labels,
minibatch_size=minibatch_gpu_in,
**cD.loss_args)
for cN in [cG, cD]:
set_optimizer_ops(cN, lazy_regularization, no_op)
# Setup training ops
data_fetch_op = tf.group(*data_fetch_ops)
for cN in [cG, cD]:
cN.train_op = cN.opt.apply_updates()
cN.reg_op = cN.reg_opt.apply_updates(allow_no_op=True)
Gs_update_op = Gs.setup_as_moving_average_of(G, beta=beta)
# Finalize graph
with tf.device("/gpu:0"):
try:
peak_gpu_mem_op = tf.contrib.memory_stats.MaxBytesInUse()
except tf.errors.NotFoundError:
peak_gpu_mem_op = tf.constant(0)
tflib.init_uninitialized_vars()
# Tensorboard summaries
if summarize:
misc.log("Initializing logs...", "white")
summary_log = tf.summary.FileWriter(dnnlib.make_run_dir_path())
if save_tf_graph:
summary_log.add_graph(tf.get_default_graph())
if save_weight_histograms:
G.setup_weight_histograms()
D.setup_weight_histograms()
# Initialize training
misc.log("Training for %d kimg..." % total_kimg, "white")
dnnlib.RunContext.get().update("",
cur_epoch=resume_kimg,
max_epoch=total_kimg)
maintenance_time = dnnlib.RunContext.get().get_last_update_interval()
cur_tick, running_mb_counter = -1, 0
cur_nimg = int(resume_kimg * 1000)
tick_start_nimg = cur_nimg
for cN in [cG, cD]:
cN.lossvals_agg = {
k: None
for k in ["loss", "reg", "norm", "reg_norm"]
}
cN.opt.reset_optimizer_state()
# Training loop
while cur_nimg < total_kimg * 1000:
if dnnlib.RunContext.get().should_stop():
break
# Choose training parameters and configure training ops
sched = training_schedule(sched_args,
cur_nimg=cur_nimg,
dataset=dataset)
assert sched.minibatch_size % (sched.minibatch_gpu * num_gpus) == 0
dataset.configure(sched.minibatch_gpu)
# Run training ops
feed_dict = {
lrate_in_g: sched.G_lrate,
lrate_in_d: sched.D_lrate,
minibatch_size_in: sched.minibatch_size,
minibatch_gpu_in: sched.minibatch_gpu,
step: sched.kimg
}
# Several iterations before updating training parameters
for _repeat in range(minibatch_repeats):
rounds = range(0, sched.minibatch_size,
sched.minibatch_gpu * num_gpus)
for cN in [cG, cD]:
cN.run_reg = lazy_regularization and (running_mb_counter %
cN.reg_interval == 0)
cur_nimg += sched.minibatch_size
running_mb_counter += 1
for cN in [cG, cD]:
cN.lossvals = {
k: None
for k in ["loss", "reg", "norm", "reg_norm"]
}
# Gradient accumulation
for _round in rounds:
cG.lossvals.update(
tflib.run([cG.train_op, cG.ops], feed_dict)[1])
if cG.run_reg:
_, cG.lossvals["reg_norm"] = tflib.run(
[cG.reg_op, cG.reg_norm], feed_dict)
tflib.run(data_fetch_op, feed_dict)
cD.lossvals.update(
tflib.run([cD.train_op, cD.ops], feed_dict)[1])
if cD.run_reg:
_, cD.lossvals["reg_norm"] = tflib.run(
[cD.reg_op, cD.reg_norm], feed_dict)
tflib.run([Gs_update_op], feed_dict)
# Track loss statistics
for cN in [cG, cD]:
for k in cN.lossvals_agg:
cN.lossvals_agg[k] = emaAvg(cN.lossvals_agg[k],
cN.lossvals[k])
# Perform maintenance tasks once per tick
done = (cur_nimg >= total_kimg * 1000)
if cur_tick < 0 or cur_nimg >= tick_start_nimg + sched.tick_kimg * 1000 or done:
cur_tick += 1
tick_kimg = (cur_nimg - tick_start_nimg) / 1000.0
tick_start_nimg = cur_nimg
tick_time = dnnlib.RunContext.get().get_time_since_last_update()
total_time = dnnlib.RunContext.get().get_time_since_start(
) + resume_time
# Report progress
print(
("tick %s kimg %s loss/reg: G (%s %s) D (%s %s) grad norms: G (%s %s) D (%s %s) "
+ "time %s sec/kimg %s maxGPU %sGB %s") %
(misc.bold("%-5d" % autosummary("Progress/tick", cur_tick)),
misc.bcolored(
"{:>8.1f}".format(
autosummary("Progress/kimg", cur_nimg / 1000.0)),
"red"),
misc.bcolored("{:>6.3f}".format(cG.lossvals_agg["loss"] or 0),
"blue"),
misc.bold("{:>6.3f}".format(cG.lossvals_agg["reg"] or 0)),
misc.bcolored("{:>6.3f}".format(cD.lossvals_agg["loss"] or 0),
"blue"),
misc.bold("{:>6.3f}".format(cD.lossvals_agg["reg"] or 0)),
misc.cond_bcolored(cG.lossvals_agg["norm"], 20.0, "red"),
misc.cond_bcolored(cG.lossvals_agg["reg_norm"], 20.0, "red"),
misc.cond_bcolored(cD.lossvals_agg["norm"], 20.0, "red"),
misc.cond_bcolored(cD.lossvals_agg["reg_norm"], 20.0, "red"),
misc.bold("%-10s" % dnnlib.util.format_time(
autosummary("Timing/total_sec", total_time))),
"{:>7.2f}".format(
autosummary("Timing/sec_per_kimg",
tick_time / tick_kimg)),
"{:>4.1f}".format(
autosummary("Resources/peak_gpu_mem_gb",
peak_gpu_mem_op.eval() / 2**30)), printname))
autosummary("Timing/total_hours", total_time / (60.0 * 60.0))
autosummary("Timing/total_days", total_time / (24.0 * 60.0 * 60.0))
# Save snapshots
if img_snapshot_ticks is not None and (
cur_tick % img_snapshot_ticks == 0 or done):
visualize.eval(Gs,
dataset,
batch_size=sched.minibatch_gpu,
training=True,
step=cur_nimg // 1000,
grid_size=grid_size,
latents=grid_latents,
labels=grid_labels,
drange_net=drange_net,
ratio=ratio,
**vis_args)
if network_snapshot_ticks is not None and (
cur_tick % network_snapshot_ticks == 0 or done):
pkl = dnnlib.make_run_dir_path("network-snapshot-%06d.pkl" %
(cur_nimg // 1000))
misc.save_pkl((G, D, Gs), pkl, remove=False)
if cur_tick % network_snapshot_ticks == 0 or done:
metric = metrics.run(
pkl,
num_imgs=eval_images_num,
run_dir=dnnlib.make_run_dir_path(),
data_dir=dnnlib.convert_path(data_dir),
num_gpus=num_gpus,
ratio=ratio,
tf_config=tf_config,
mirror_augment=mirror_augment)
if last_snapshots > 0:
misc.rm(
sorted(
glob.glob(dnnlib.make_run_dir_path(
"network*.pkl")))[:-last_snapshots])
# Update summaries and RunContext
if summarize:
metrics.update_autosummaries()
tflib.autosummary.save_summaries(summary_log, cur_nimg)
dnnlib.RunContext.get().update(None,
cur_epoch=cur_nimg // 1000,
max_epoch=total_kimg)
maintenance_time = dnnlib.RunContext.get(
).get_last_update_interval() - tick_time
# Save final snapshot
misc.save_pkl((G, D, Gs),
dnnlib.make_run_dir_path("network-final.pkl"),
remove=False)
# All done
if summarize:
summary_log.close()
dataset.close()
def training_loop(
run_dir='.', # Output directory.
G_args={}, # Options for generator network.
D_args={}, # Options for discriminator network.
G_opt_args={}, # Options for generator optimizer.
D_opt_args={}, # Options for discriminator optimizer.
loss_args={}, # Options for loss function.
train_dataset_args={}, # Options for dataset to train with.
metric_dataset_args={}, # Options for dataset to evaluate metrics against.
augment_args={}, # Options for adaptive augmentations.
metric_arg_list=[], # Metrics to evaluate during training.
num_gpus=1, # Number of GPUs to use.
minibatch_size=32, # Global minibatch size.
minibatch_gpu=4, # Number of samples processed at a time by one GPU.
G_smoothing_kimg=10, # Half-life of the exponential moving average (EMA) of generator weights.
G_smoothing_rampup=None, # EMA ramp-up coefficient.
minibatch_repeats=4, # Number of minibatches to run in the inner loop.
lazy_regularization=True, # Perform regularization as a separate training step?
G_reg_interval=4, # How often the perform regularization for G? Ignored if lazy_regularization=False.
D_reg_interval=16, # How often the perform regularization for D? Ignored if lazy_regularization=False.
total_kimg=25000, # Total length of the training, measured in thousands of real images.
kimg_per_tick=10, # Progress snapshot interval.
image_snapshot_ticks=1, # How often to save image snapshots? None = only save 'reals.png' and 'fakes-init.png'.
network_snapshot_ticks=1, # How often to save network snapshots? None = only save 'networks-final.pkl'.
resume_pkl=None, # Network pickle to resume training from, None = train from scratch.
resume_kimg=15000, # Assumed training progress at the beginning. Affects reporting and training schedule.
resume_time=0.0, # Assumed wallclock time at the beginning. Affects reporting.
abort_fn=None, # Callback function for determining whether to abort training.
progress_fn=None, # Callback function for updating training progress.
):
assert minibatch_size % (num_gpus * minibatch_gpu) == 0
start_time = time.time()
print('Loading training set...')
training_set = dataset.load_dataset(**train_dataset_args)
print('Image shape:', np.int32(training_set.shape).tolist())
print('Label shape:', [training_set.label_size])
print()
print('Constructing networks...')
with tf.device('/gpu:0'):
G = tflib.Network('G',
num_channels=training_set.shape[0],
resolution=training_set.shape[1],
label_size=training_set.label_size,
**G_args)
D = tflib.Network('D',
num_channels=training_set.shape[0],
resolution=training_set.shape[1],
label_size=training_set.label_size,
**D_args)
Gs = G.clone('Gs')
if resume_pkl is not None:
print(f'Resuming from "{resume_pkl}"')
with dnnlib.util.open_url(resume_pkl) as f:
#[EDITED]
G, D, Gs = pickle.load(f)
# rG, rD, rGs = pickle.load(f)
# G.copy_vars_from(rG)
# D.copy_vars_from(rD)
# Gs.copy_vars_from(rGs)
# G.print_layers()
# D.print_layers()
start_time = time.time()
grid_size, grid_reals, grid_labels = setup_snapshot_image_grid(
training_set)
end_time = time.time()
print('Finish setting up in ', end_time - start_time)
start_time = time.time()
save_image_grid(grid_reals,
os.path.join(run_dir, 'reals.png'),
drange=[0, 255],
grid_size=grid_size)
end_time = time.time()
print('Finished saving in ', end_time - start_time)
grid_latents = np.random.randn(np.prod(grid_size), *G.input_shape[1:])
grid_fakes = Gs.run(grid_latents,
grid_labels,
is_validation=True,
minibatch_size=minibatch_gpu)
save_image_grid(grid_fakes,
os.path.join(run_dir, 'fakes_init.png'),
drange=[-1, 1],
grid_size=grid_size)
print(f'Replicating networks across {num_gpus} GPUs...')
G_gpus = [G]
D_gpus = [D]
for gpu in range(1, num_gpus):
with tf.device(f'/gpu:{gpu}'):
G_gpus.append(G.clone(f'{G.name}_gpu{gpu}'))
D_gpus.append(D.clone(f'{D.name}_gpu{gpu}'))
print('Initializing augmentations...')
aug = None
if augment_args.get('class_name', None) is not None:
aug = dnnlib.util.construct_class_by_name(**augment_args)
aug.init_validation_set(D_gpus=D_gpus, training_set=training_set)
print('Setting up optimizers...')
G_opt_args = dict(G_opt_args)
D_opt_args = dict(D_opt_args)
for args, reg_interval in [(G_opt_args, G_reg_interval),
(D_opt_args, D_reg_interval)]:
args[
'minibatch_multiplier'] = minibatch_size // num_gpus // minibatch_gpu
if lazy_regularization:
mb_ratio = reg_interval / (reg_interval + 1)
args['learning_rate'] *= mb_ratio
if 'beta1' in args: args['beta1'] **= mb_ratio
if 'beta2' in args: args['beta2'] **= mb_ratio
G_opt = tflib.Optimizer(name='TrainG', **G_opt_args)
D_opt = tflib.Optimizer(name='TrainD', **D_opt_args)
G_reg_opt = tflib.Optimizer(name='RegG', share=G_opt, **G_opt_args)
D_reg_opt = tflib.Optimizer(name='RegD', share=D_opt, **D_opt_args)
print('Constructing training graph...')
data_fetch_ops = []
training_set.configure(minibatch_gpu)
for gpu, (G_gpu, D_gpu) in enumerate(zip(G_gpus, D_gpus)):
with tf.name_scope(f'Train_gpu{gpu}'), tf.device(f'/gpu:{gpu}'):
# Fetch training data via temporary variables.
with tf.name_scope('DataFetch'):
real_images_var = tf.Variable(
name='images',
trainable=False,
initial_value=tf.zeros([minibatch_gpu] +
training_set.shape))
real_labels_var = tf.Variable(name='labels',
trainable=False,
initial_value=tf.zeros([
minibatch_gpu,
training_set.label_size
]))
real_images_write, real_labels_write = training_set.get_minibatch_tf(
)
real_images_write = tflib.convert_images_from_uint8(
real_images_write)
data_fetch_ops += [
tf.assign(real_images_var, real_images_write)
]
data_fetch_ops += [
tf.assign(real_labels_var, real_labels_write)
]
# Evaluate loss function and register gradients.
fake_labels = training_set.get_random_labels_tf(minibatch_gpu)
terms = dnnlib.util.call_func_by_name(G=G_gpu,
D=D_gpu,
aug=aug,
fake_labels=fake_labels,
real_images=real_images_var,
real_labels=real_labels_var,
**loss_args)
if lazy_regularization:
if terms.G_reg is not None:
G_reg_opt.register_gradients(
tf.reduce_mean(terms.G_reg * G_reg_interval),
G_gpu.trainables)
if terms.D_reg is not None:
D_reg_opt.register_gradients(
tf.reduce_mean(terms.D_reg * D_reg_interval),
D_gpu.trainables)
else:
if terms.G_reg is not None: terms.G_loss += terms.G_reg
if terms.D_reg is not None: terms.D_loss += terms.D_reg
G_opt.register_gradients(tf.reduce_mean(terms.G_loss),
G_gpu.trainables)
D_opt.register_gradients(tf.reduce_mean(terms.D_loss),
D_gpu.trainables)
print('Finalizing training ops...')
data_fetch_op = tf.group(*data_fetch_ops)
G_train_op = G_opt.apply_updates()
D_train_op = D_opt.apply_updates()
G_reg_op = G_reg_opt.apply_updates(allow_no_op=True)
D_reg_op = D_reg_opt.apply_updates(allow_no_op=True)
Gs_beta_in = tf.placeholder(tf.float32, name='Gs_beta_in', shape=[])
Gs_update_op = Gs.setup_as_moving_average_of(G, beta=Gs_beta_in)
tflib.init_uninitialized_vars()
with tf.device('/gpu:0'):
peak_gpu_mem_op = tf.contrib.memory_stats.MaxBytesInUse()
print('Initializing metrics...')
summary_log = tf.summary.FileWriter(run_dir)
metrics = []
for args in metric_arg_list:
metric = dnnlib.util.construct_class_by_name(**args)
metric.configure(dataset_args=metric_dataset_args, run_dir=run_dir)
metrics.append(metric)
print(f'Training for {total_kimg} kimg...')
print()
if progress_fn is not None:
progress_fn(0, total_kimg)
tick_start_time = time.time()
maintenance_time = tick_start_time - start_time
cur_nimg = 0
cur_tick = -1
tick_start_nimg = cur_nimg
running_mb_counter = 0
done = False
while not done:
# Compute EMA decay parameter.
Gs_nimg = G_smoothing_kimg * 1000.0
if G_smoothing_rampup is not None:
Gs_nimg = min(Gs_nimg, cur_nimg * G_smoothing_rampup)
Gs_beta = 0.5**(minibatch_size / max(Gs_nimg, 1e-8))
# Run training ops.
for _repeat_idx in range(minibatch_repeats):
rounds = range(0, minibatch_size, minibatch_gpu * num_gpus)
run_G_reg = (lazy_regularization
and running_mb_counter % G_reg_interval == 0)
run_D_reg = (lazy_regularization
and running_mb_counter % D_reg_interval == 0)
cur_nimg += minibatch_size
running_mb_counter += 1
# Fast path without gradient accumulation.
if len(rounds) == 1:
tflib.run([G_train_op, data_fetch_op])
if run_G_reg:
tflib.run(G_reg_op)
tflib.run([D_train_op, Gs_update_op], {Gs_beta_in: Gs_beta})
if run_D_reg:
tflib.run(D_reg_op)
# Slow path with gradient accumulation.
else:
for _round in rounds:
tflib.run(G_train_op)
if run_G_reg:
tflib.run(G_reg_op)
tflib.run(Gs_update_op, {Gs_beta_in: Gs_beta})
for _round in rounds:
tflib.run(data_fetch_op)
tflib.run(D_train_op)
if run_D_reg:
tflib.run(D_reg_op)
# Run validation.
if aug is not None:
aug.run_validation(minibatch_size=minibatch_size)
# Tune augmentation parameters.
if aug is not None:
aug.tune(minibatch_size * minibatch_repeats)
# Perform maintenance tasks once per tick.
done = (cur_nimg >= total_kimg * 1000) or (abort_fn is not None
and abort_fn())
if done or cur_tick < 0 or cur_nimg >= tick_start_nimg + kimg_per_tick * 1000:
cur_tick += 1
tick_kimg = (cur_nimg - tick_start_nimg) / 1000.0
tick_start_nimg = cur_nimg
tick_end_time = time.time()
total_time = tick_end_time - start_time
tick_time = tick_end_time - tick_start_time
# Report progress.
print(' '.join([
f"tick {autosummary('Progress/tick', cur_tick):<5d}",
f"kimg {autosummary('Progress/kimg', cur_nimg / 1000.0):<8.1f}",
f"time {dnnlib.util.format_time(autosummary('Timing/total_sec', total_time)):<12s}",
f"sec/tick {autosummary('Timing/sec_per_tick', tick_time):<7.1f}",
f"sec/kimg {autosummary('Timing/sec_per_kimg', tick_time / tick_kimg):<7.2f}",
f"maintenance {autosummary('Timing/maintenance_sec', maintenance_time):<6.1f}",
f"gpumem {autosummary('Resources/peak_gpu_mem_gb', peak_gpu_mem_op.eval() / 2**30):<5.1f}",
f"augment {autosummary('Progress/augment', aug.strength if aug is not None else 0):.3f}",
]))
autosummary('Timing/total_hours', total_time / (60.0 * 60.0))
autosummary('Timing/total_days', total_time / (24.0 * 60.0 * 60.0))
if progress_fn is not None:
progress_fn(cur_nimg // 1000, total_kimg)
# Save snapshots.
if image_snapshot_ticks is not None and (
done or cur_tick % image_snapshot_ticks == 0):
grid_fakes = Gs.run(grid_latents,
grid_labels,
is_validation=True,
minibatch_size=minibatch_gpu)
save_image_grid(grid_fakes,
os.path.join(
run_dir,
f'fakes{cur_nimg // 1000:06d}.png'),
drange=[-1, 1],
grid_size=grid_size)
if network_snapshot_ticks is not None and (
done or cur_tick % network_snapshot_ticks == 0):
pkl = os.path.join(
run_dir, f'network-snapshot-{cur_nimg // 1000:06d}.pkl')
with open(pkl, 'wb') as f:
pickle.dump((G, D, Gs), f)
if len(metrics):
print('Evaluating metrics...')
for metric in metrics:
metric.run(pkl, num_gpus=num_gpus)
# Update summaries.
for metric in metrics:
metric.update_autosummaries()
tflib.autosummary.save_summaries(summary_log, cur_nimg)
tick_start_time = time.time()
maintenance_time = tick_start_time - tick_end_time
print()
print('Exiting...')
summary_log.close()
training_set.close()
def get_random_labels_np(self, minibatch_size): # => labels
self.configure(minibatch_size)
if self._tf_labels_np is None:
self._tf_labels_np = self.get_random_labels_tf(minibatch_size)
return tflib.run(self._tf_labels_np)
def training_loop(
G_args={}, # Options for generator network.
D_args={}, # Options for discriminator network.
G_opt_args={}, # Options for generator optimizer.
D_opt_args={}, # Options for discriminator optimizer.
G_loss_args={}, # Options for generator loss.
D_loss_args={}, # Options for discriminator loss.
dataset_args={}, # Options for dataset.load_dataset().
sched_args={}, # Options for train.TrainingSchedule.
grid_args={}, # Options for train.setup_snapshot_image_grid().
metric_arg_list=[], # Options for MetricGroup.
tf_config={}, # Options for tflib.init_tf().
data_dir=None, # Directory to load datasets from.
G_smoothing_kimg=10.0, # Half-life of the running average of generator weights.
minibatch_repeats=4, # Number of minibatches to run before adjusting training parameters.
lazy_regularization=True, # Perform regularization as a separate training step?
G_reg_interval=4, # How often the perform regularization for G? Ignored if lazy_regularization=False.
D_reg_interval=16, # How often the perform regularization for D? Ignored if lazy_regularization=False.
reset_opt_for_new_lod=True, # Reset optimizer internal state (e.g. Adam moments) when new layers are introduced?
total_kimg=25000, # Total length of the training, measured in thousands of real images.
mirror_augment=False, # Enable mirror augment?
drange_net=[
-1, 1
], # Dynamic range used when feeding image data to the networks.
image_snapshot_ticks=50, # How often to save image snapshots? None = only save 'reals.png' and 'fakes-init.png'.
network_snapshot_ticks=50, # How often to save network snapshots? None = only save 'networks-final.pkl'.
save_tf_graph=False, # Include full TensorFlow computation graph in the tfevents file?
save_weight_histograms=False, # Include weight histograms in the tfevents file?
resume_pkl=None, # Network pickle to resume training from, None = train from scratch.
resume_kimg=0.0, # Assumed training progress at the beginning. Affects reporting and training schedule.
resume_time=0.0, # Assumed wallclock time at the beginning. Affects reporting.
resume_with_new_nets=False
): # Construct new networks according to G_args and D_args before resuming training?
# Initialize dnnlib and TensorFlow.
tflib.init_tf(tf_config)
num_gpus = dnnlib.submit_config.num_gpus
# Load training set.
training_set = dataset.load_dataset(data_dir=dnnlib.convert_path(data_dir),
verbose=True,
**dataset_args)
training_set.label_size += 1
grid_size, grid_reals, grid_labels = misc.setup_snapshot_image_grid(
training_set, mirror_label=True, **grid_args)
misc.save_image_grid(grid_reals,
dnnlib.make_run_dir_path('reals.png'),
drange=training_set.dynamic_range,
grid_size=grid_size)
# Construct or load networks.
with tf.device('/gpu:0'):
if resume_pkl is None or resume_with_new_nets:
print('Constructing networks...')
G = tflib.Network('G',
num_channels=training_set.shape[0],
resolution=training_set.shape[1],
label_size=training_set.label_size,
**G_args)
D = tflib.Network('D',
num_channels=training_set.shape[0],
resolution=training_set.shape[1],
label_size=training_set.label_size,
**D_args)
Gs = G.clone('Gs')
if resume_pkl is not None:
print('Loading networks from "%s"...' % resume_pkl)
rG, rD, rGs = misc.load_pkl(resume_pkl)
if resume_with_new_nets:
G.copy_vars_from(rG)
D.copy_vars_from(rD)
Gs.copy_vars_from(rGs)
else:
G = rG
D = rD
Gs = rGs
# Print layers and generate initial image snapshot.
G.print_layers()
D.print_layers()
sched = training_schedule(cur_nimg=total_kimg * 1000,
training_set=training_set,
**sched_args)
grid_latents = np.random.randn(np.prod(grid_size), *G.input_shape[1:])
# Setup training inputs.
print('Building TensorFlow graph...')
with tf.name_scope('Inputs'), tf.device('/cpu:0'):
lod_in = tf.placeholder(tf.float32, name='lod_in', shape=[])
lrate_in = tf.placeholder(tf.float32, name='lrate_in', shape=[])
minibatch_size_in = tf.placeholder(tf.int32,
name='minibatch_size_in',
shape=[])
minibatch_gpu_in = tf.placeholder(tf.int32,
name='minibatch_gpu_in',
shape=[])
minibatch_multiplier = minibatch_size_in // (minibatch_gpu_in *
num_gpus)
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
# Setup optimizers.
G_opt_args = dict(G_opt_args)
D_opt_args = dict(D_opt_args)
for args, reg_interval in [(G_opt_args, G_reg_interval),
(D_opt_args, D_reg_interval)]:
args['minibatch_multiplier'] = minibatch_multiplier
args['learning_rate'] = lrate_in
if lazy_regularization:
mb_ratio = reg_interval / (reg_interval + 1)
args['learning_rate'] *= mb_ratio
if 'beta1' in args: args['beta1'] **= mb_ratio
if 'beta2' in args: args['beta2'] **= mb_ratio
G_opt = tflib.Optimizer(name='TrainG', **G_opt_args)
D_opt = tflib.Optimizer(name='TrainD', **D_opt_args)
G_reg_opt = tflib.Optimizer(name='RegG', share=G_opt, **G_opt_args)
D_reg_opt = tflib.Optimizer(name='RegD', share=D_opt, **D_opt_args)
# Build training graph for each GPU.
data_fetch_ops = []
for gpu in range(num_gpus):
with tf.name_scope('GPU%d' % gpu), tf.device('/gpu:%d' % gpu):
# Create GPU-specific shadow copies of G and D.
G_gpu = G if gpu == 0 else G.clone(G.name + '_shadow')
D_gpu = D if gpu == 0 else D.clone(D.name + '_shadow')
# Fetch training data via temporary variables.
with tf.name_scope('DataFetch'):
sched = training_schedule(cur_nimg=int(resume_kimg * 1000),
training_set=training_set,
**sched_args)
reals_var = tf.Variable(
name='reals',
trainable=False,
initial_value=tf.zeros([sched.minibatch_gpu] +
training_set.shape))
labels_var = tf.Variable(name='labels',
trainable=False,
initial_value=tf.zeros([
sched.minibatch_gpu,
training_set.label_size
]))
reals_write, labels_write = training_set.get_minibatch_tf()
reals_write, labels_write = process_reals(
reals_write, labels_write, lod_in, mirror_augment,
training_set.dynamic_range, drange_net)
reals_write = tf.concat(
[reals_write, reals_var[minibatch_gpu_in:]], axis=0)
labels_write = tf.concat(
[labels_write, labels_var[minibatch_gpu_in:]], axis=0)
data_fetch_ops += [tf.assign(reals_var, reals_write)]
data_fetch_ops += [tf.assign(labels_var, labels_write)]
reals_read = reals_var[:minibatch_gpu_in]
labels_read = labels_var[:minibatch_gpu_in]
# Evaluate loss functions.
lod_assign_ops = []
if 'lod' in G_gpu.vars:
lod_assign_ops += [tf.assign(G_gpu.vars['lod'], lod_in)]
if 'lod' in D_gpu.vars:
lod_assign_ops += [tf.assign(D_gpu.vars['lod'], lod_in)]
with tf.control_dependencies(lod_assign_ops):
with tf.name_scope('G_loss'):
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)
with tf.name_scope('D_loss'):
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)
# Register gradients.
if not lazy_regularization:
if G_reg is not None: G_loss += G_reg
if D_reg is not None: D_loss += D_reg
else:
if G_reg is not None:
G_reg_opt.register_gradients(
tf.reduce_mean(G_reg * G_reg_interval),
G_gpu.trainables)
if D_reg is not None:
D_reg_opt.register_gradients(
tf.reduce_mean(D_reg * D_reg_interval),
D_gpu.trainables)
G_opt.register_gradients(tf.reduce_mean(G_loss), G_gpu.trainables)
D_opt.register_gradients(tf.reduce_mean(D_loss), D_gpu.trainables)
# Setup training ops.
data_fetch_op = tf.group(*data_fetch_ops)
G_train_op = G_opt.apply_updates()
D_train_op = D_opt.apply_updates()
G_reg_op = G_reg_opt.apply_updates(allow_no_op=True)
D_reg_op = D_reg_opt.apply_updates(allow_no_op=True)
Gs_update_op = Gs.setup_as_moving_average_of(G, beta=Gs_beta)
# Finalize graph.
with tf.device('/gpu:0'):
try:
peak_gpu_mem_op = tf.contrib.memory_stats.MaxBytesInUse()
except tf.errors.NotFoundError:
peak_gpu_mem_op = tf.constant(0)
tflib.init_uninitialized_vars()
print('Initializing logs...')
summary_log = tf.summary.FileWriter(dnnlib.make_run_dir_path())
if save_tf_graph:
summary_log.add_graph(tf.get_default_graph())
if save_weight_histograms:
G.setup_weight_histograms()
D.setup_weight_histograms()
metrics = metric_base.MetricGroup(metric_arg_list)
print('Training for %d kimg...\n' % total_kimg)
dnnlib.RunContext.get().update('',
cur_epoch=resume_kimg,
max_epoch=total_kimg)
maintenance_time = dnnlib.RunContext.get().get_last_update_interval()
cur_nimg = int(resume_kimg * 1000)
cur_tick = -1
tick_start_nimg = cur_nimg
prev_lod = -1.0
running_mb_counter = 0
while cur_nimg < total_kimg * 1000:
if dnnlib.RunContext.get().should_stop(): break
# Choose training parameters and configure training ops.
sched = training_schedule(cur_nimg=cur_nimg,
training_set=training_set,
**sched_args)
assert sched.minibatch_size % (sched.minibatch_gpu * num_gpus) == 0
training_set.configure(sched.minibatch_gpu, sched.lod)
if reset_opt_for_new_lod:
if np.floor(sched.lod) != np.floor(prev_lod) or np.ceil(
sched.lod) != np.ceil(prev_lod):
G_opt.reset_optimizer_state()
D_opt.reset_optimizer_state()
prev_lod = sched.lod
# Run training ops.
feed_dict = {
lod_in: sched.lod,
lrate_in: sched.G_lrate,
minibatch_size_in: sched.minibatch_size,
minibatch_gpu_in: sched.minibatch_gpu
}
for _repeat in range(minibatch_repeats):
rounds = range(0, sched.minibatch_size,
sched.minibatch_gpu * num_gpus)
run_G_reg = (lazy_regularization
and running_mb_counter % G_reg_interval == 0)
run_D_reg = (lazy_regularization
and running_mb_counter % D_reg_interval == 0)
cur_nimg += sched.minibatch_size
running_mb_counter += 1
# Fast path without gradient accumulation.
if len(rounds) == 1:
tflib.run([G_train_op, data_fetch_op], feed_dict)
if run_G_reg:
tflib.run(G_reg_op, feed_dict)
tflib.run([D_train_op, Gs_update_op], feed_dict)
if run_D_reg:
tflib.run(D_reg_op, feed_dict)
# Slow path with gradient accumulation.
else:
for _round in rounds:
tflib.run(G_train_op, feed_dict)
if run_G_reg:
for _round in rounds:
tflib.run(G_reg_op, feed_dict)
tflib.run(Gs_update_op, feed_dict)
for _round in rounds:
tflib.run(data_fetch_op, feed_dict)
tflib.run(D_train_op, feed_dict)
if run_D_reg:
for _round in rounds:
tflib.run(D_reg_op, feed_dict)
# Perform maintenance tasks once per tick.
done = (cur_nimg >= total_kimg * 1000)
if cur_tick < 0 or cur_nimg >= tick_start_nimg + sched.tick_kimg * 1000 or done:
cur_tick += 1
tick_kimg = (cur_nimg - tick_start_nimg) / 1000.0
tick_start_nimg = cur_nimg
tick_time = dnnlib.RunContext.get().get_time_since_last_update()
total_time = dnnlib.RunContext.get().get_time_since_start(
) + resume_time
# Report progress.
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'
% (autosummary('Progress/tick', cur_tick),
autosummary('Progress/kimg', cur_nimg / 1000.0),
autosummary('Progress/lod', sched.lod),
autosummary('Progress/minibatch', sched.minibatch_size),
dnnlib.util.format_time(
autosummary('Timing/total_sec', total_time)),
autosummary('Timing/sec_per_tick', tick_time),
autosummary('Timing/sec_per_kimg', tick_time / tick_kimg),
autosummary('Timing/maintenance_sec', maintenance_time),
autosummary('Resources/peak_gpu_mem_gb',
peak_gpu_mem_op.eval() / 2**30)))
autosummary('Timing/total_hours', total_time / (60.0 * 60.0))
autosummary('Timing/total_days', total_time / (24.0 * 60.0 * 60.0))
# Save snapshots.
if image_snapshot_ticks is not None and (
cur_tick % image_snapshot_ticks == 0 or done):
grid_fakes = Gs.run(grid_latents,
grid_labels,
is_validation=True,
minibatch_size=sched.minibatch_gpu)
print(grid_labels.shape)
misc.save_image_grid(grid_fakes,
dnnlib.make_run_dir_path(
'fakes%06d.png' % (cur_nimg // 1000)),
drange=drange_net,
grid_size=grid_size)
# generate_images_with_labels(dnnlib.make_run_dir_path('labels%06d.png' % (cur_nimg // 1000)), Gs, w=256, h=256, num_labels=training_set.label_size, latents=label_grid_latents)
if network_snapshot_ticks is not None and (
cur_tick % network_snapshot_ticks == 0 or done):
pkl = dnnlib.make_run_dir_path('network-snapshot-%06d.pkl' %
(cur_nimg // 1000))
misc.save_pkl((G, D, Gs), pkl)
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)
# Update summaries and RunContext.
metrics.update_autosummaries()
tflib.autosummary.save_summaries(summary_log, cur_nimg)
dnnlib.RunContext.get().update('%.2f' % sched.lod,
cur_epoch=cur_nimg // 1000,
max_epoch=total_kimg)
maintenance_time = dnnlib.RunContext.get(
).get_last_update_interval() - tick_time
# Save final snapshot.
misc.save_pkl((G, D, Gs), dnnlib.make_run_dir_path('network-final.pkl'))
# All done.
summary_log.close()
training_set.close()
def dlatents(self):
return tflib.run(self._dlatents_expr, {self._dlatent_noise_in: 0})
def _evaluate(self, Gs, Gs_kwargs, num_gpus):
Gs_kwargs = dict(Gs_kwargs)
Gs_kwargs.update(self.Gs_overrides)
minibatch_size = num_gpus * self.minibatch_per_gpu
# Construct TensorFlow graph.
distance_expr = []
for gpu_idx in range(num_gpus):
with tf.device('/gpu:%d' % gpu_idx):
Gs_clone = Gs.clone()
noise_vars = [
var for name, var in
Gs_clone.components.synthesis.vars.items()
if name.startswith('noise')
]
# Generate random latents and interpolation t-values.
lat_t01 = tf.random_normal([self.minibatch_per_gpu * 2] +
Gs_clone.input_shape[1:])
lerp_t = tf.random_uniform(
[self.minibatch_per_gpu], 0.0,
1.0 if self.sampling == 'full' else 0.0)
labels = tf.reshape(
tf.tile(self._get_random_labels_tf(self.minibatch_per_gpu),
[1, 2]), [self.minibatch_per_gpu * 2, -1])
# Interpolate in W or Z.
if self.space == 'w':
dlat_t01 = Gs_clone.components.mapping.get_output_for(
lat_t01, labels, **Gs_kwargs)
dlat_t01 = tf.cast(dlat_t01, tf.float32)
dlat_t0, dlat_t1 = dlat_t01[0::2], dlat_t01[1::2]
dlat_e0 = tflib.lerp(dlat_t0, dlat_t1,
lerp_t[:, np.newaxis, np.newaxis])
dlat_e1 = tflib.lerp(
dlat_t0, dlat_t1,
lerp_t[:, np.newaxis, np.newaxis] + self.epsilon)
dlat_e01 = tf.reshape(tf.stack([dlat_e0, dlat_e1], axis=1),
dlat_t01.shape)
else: # space == 'z'
lat_t0, lat_t1 = lat_t01[0::2], lat_t01[1::2]
lat_e0 = slerp(lat_t0, lat_t1, lerp_t[:, np.newaxis])
lat_e1 = slerp(lat_t0, lat_t1,
lerp_t[:, np.newaxis] + self.epsilon)
lat_e01 = tf.reshape(tf.stack([lat_e0, lat_e1], axis=1),
lat_t01.shape)
dlat_e01 = Gs_clone.components.mapping.get_output_for(
lat_e01, labels, **Gs_kwargs)
# Synthesize images.
with tf.control_dependencies([
var.initializer for var in noise_vars
]): # use same noise inputs for the entire minibatch
images = Gs_clone.components.synthesis.get_output_for(
dlat_e01, randomize_noise=False, **Gs_kwargs)
images = tf.cast(images, tf.float32)
# Crop only the face region.
if self.crop:
c = int(images.shape[2] // 8)
images = images[:, :, c * 3:c * 7, c * 2:c * 6]
# Downsample image to 256x256 if it's larger than that. VGG was built for 224x224 images.
factor = images.shape[2] // 256
if factor > 1:
images = tf.reshape(images, [
-1, images.shape[1], images.shape[2] // factor, factor,
images.shape[3] // factor, factor
])
images = tf.reduce_mean(images, axis=[3, 5])
# Scale dynamic range from [-1,1] to [0,255] for VGG.
images = (images + 1) * (255 / 2)
# Evaluate perceptual distance.
img_e0, img_e1 = images[0::2], images[1::2]
distance_measure = misc.load_pkl(
'https://s3.eu-central-1.wasabisys.com/sheerun/models/vgg16_zhang_perceptual.pkl'
) # vgg16_zhang_perceptual.pkl
distance_expr.append(
distance_measure.get_output_for(img_e0, img_e1) *
(1 / self.epsilon**2))
# Sampling loop.
all_distances = []
for begin in range(0, self.num_samples, minibatch_size):
self._report_progress(begin, self.num_samples)
all_distances += tflib.run(distance_expr)
all_distances = np.concatenate(all_distances, axis=0)
# Reject outliers.
lo = np.percentile(all_distances, 1, interpolation='lower')
hi = np.percentile(all_distances, 99, interpolation='higher')
filtered_distances = np.extract(
np.logical_and(lo <= all_distances, all_distances <= hi),
all_distances)
self._report_result(np.mean(filtered_distances))
def images_float(self):
return tflib.run(self._images_float_expr, {self._dlatent_noise_in: 0})
def _evaluate(self, Gs, Gs_kwargs, num_gpus):
minibatch_size = num_gpus * self.minibatch_per_gpu
# Construct TensorFlow graph.
distance_expr = []
for gpu_idx in range(num_gpus):
with tf.device('/gpu:%d' % gpu_idx):
Gs_clone = Gs.clone()
latents = tf.random_normal([self.minibatch_per_gpu] +
Gs_clone.input_shape[1:])
labels = self._get_random_labels_tf(self.minibatch_per_gpu)
rotation_offset = 108
# replace background with black or white
if self.only_black_white_background:
background_offset = 1 + 67 + 12 + 18 + 10 + 8 + 5
all_rotations = tf.constant(
[[1, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 1]],
dtype=tf.float32)
indices = tf.cast(tf.floor(
tf.random_uniform(shape=[self.minibatch_per_gpu],
minval=0,
maxval=2)),
dtype=tf.int32)
random_background = tf.gather(all_rotations, indices)
labels = tf.concat(
[labels[:, :background_offset], random_background],
axis=-1)
random_index = tf.cast(tf.floor(
tf.random_uniform([self.minibatch_per_gpu],
minval=0,
maxval=8)),
dtype=tf.int32)
random_one_hot = tf.one_hot(random_index, depth=8)
label_left = tf.concat([
labels[:, :rotation_offset], random_one_hot,
labels[:, rotation_offset + 8:]
],
axis=-1)
random_index_shift = tf.cast(tf.mod(random_index + 1, 8),
dtype=tf.int32)
random_one_hot_shift = tf.one_hot(random_index_shift, depth=8)
label_right = tf.concat([
labels[:, :rotation_offset], random_one_hot_shift,
labels[:, rotation_offset + 8:]
],
axis=-1)
magnitude = tf.random.uniform(
shape=[self.minibatch_per_gpu, 1], minval=0, maxval=1)
label_int_left = label_left * (
1 - magnitude) + label_right * magnitude
label_int_right = label_left * (
1 - (magnitude + self.angle)) + label_right * (magnitude +
self.angle)
images_left = Gs_clone.get_output_for(latents,
label_int_left,
randomize_noise=False)
images_left = tf.cast(images_left, tf.float32)
images_right = Gs_clone.get_output_for(latents,
label_int_right,
randomize_noise=False)
images_right = tf.cast(images_right, tf.float32)
# Downsample image to 256x256 if it's larger than that. VGG was built for 224x224 images.
# factor = images.shape[2] // 256
# if factor > 1:
# images = tf.reshape(images, [-1, images.shape[1], images.shape[2] // factor, factor, images.shape[3] // factor, factor])
# images = tf.reduce_mean(images, axis=[3,5])
#
# Scale dynamic range from [-1,1] to [0,255] for VGG.
images_left = (images_left + 1) * (255 / 2)
images_right = (images_right + 1) * (255 / 2)
# Evaluate perceptual distance.
distance_measure = misc.load_pkl(
'https://nvlabs-fi-cdn.nvidia.com/stylegan/networks/metrics/vgg16_zhang_perceptual.pkl'
)
distance_expr.append(
distance_measure.get_output_for(images_left, images_right))
# Sampling loop.
all_distances = []
for begin in range(0, self.num_images, minibatch_size):
self._report_progress(begin, self.num_images)
all_distances += tflib.run(distance_expr)
all_distances = np.concatenate(all_distances, axis=0)
# Reject outliers.
# lo = np.percentile(all_distances, 1, interpolation='lower')
# hi = np.percentile(all_distances, 99, interpolation='higher')
# filtered_distances = np.extract(np.logical_and(lo <= all_distances, all_distances <= hi), all_distances)
self._report_result(np.mean(all_distances), suffix='_mean')
self._report_result(np.std(all_distances), suffix='_stddev')
# int_pl_lengths = tf.reduce_sum(tf.abs(pl_grads))
percept_dist = distance_measure.get_output_for(image, image_2)
latents = np.random.normal(size=[minibatch_size] + G.input_shapes[0][1:])
labels = training_set.get_random_labels_np(minibatch_size)
angles = np.arange(0, np.math.pi * 2, np.math.pi / 50)
y_gradient = []
y_percept_dist = []
images = []
for i in tqdm(range(len(angles))):
int_pl_lengths_out, image_out, percept_dist_out = tflib.run(
[int_pl_lengths, image, percept_dist],
feed_dict={
latent_placeholder: latents,
label_placeholder: labels,
angle_placeholder: np.expand_dims(angles[i], axis=-1)
})
y_gradient.append(int_pl_lengths_out)
y_percept_dist.append(percept_dist_out[0])
images.append(image_out)
num_images = 15
images = np.array(images)
canvas = Image.new('RGB', (256 * num_images, 256), 'white')
for i in range(num_images):
j = int(np.round((i / num_images * images.shape[0])))
img = np.transpose(images[j, 0], [1, 2, 0])
img = img * 127.5 + 127.5
img = np.clip(img, 0, 255).astype(np.uint8)
def start(self,
target_images,
advect_last_frame_weight=0.0,
advect_last_frame_img=None,
advect_speed_mask=None,
advect_noise_vars=[],
advect_noise_weight=0.0,
noise_speed_masks={}):
assert self._Gs is not None
# Prepare target images.
self._info('Preparing target images...')
target_images = self.downsample_raw_img(target_images)
# Initialize optimization state.
self._info('Initializing optimization state...')
tflib.set_vars({
self._target_images_var:
target_images,
self._dlatents_var:
np.tile(self._dlatent_avg, [self._minibatch_size, 1, 1])
})
if self._advect_last_frame:
if advect_last_frame_img is None:
advect_last_frame_weight = 0.0
tflib.set_vars({self._advect_last_frame_img: target_images})
else:
advect_last_frame_img = self.downsample_raw_img(
advect_last_frame_img)
tflib.set_vars(
{self._advect_last_frame_img: advect_last_frame_img})
if self._mask_speed_vector:
if advect_speed_mask is None:
advect_speed_mask = np.ones(
(target_images.shape[0], 1, target_images.shape[2],
target_images.shape[3]))
else:
advect_speed_mask = self.downsample_raw_img(
advect_speed_mask, rescale=False)
tflib.set_vars({self._speed_mask: advect_speed_mask})
tflib.set_vars(
{self._advect_last_frame_weight: advect_last_frame_weight})
if self._encourage_advected_noise:
# assert noise are all in correct shape
set_dict = {}
if len(advect_noise_vars) == 0:
advect_noise_weight = 0.0
for noise_idx in range(len(self._advect_noise_vars)):
set_dict[self._advect_noise_vars[noise_idx]] = np.zeros(
self._advect_noise_vars[noise_idx].shape)
else:
assert len(advect_noise_vars) == len(self._noise_vars)
for noise_idx in range(len(self._advect_noise_vars)):
set_dict[self._advect_noise_vars[
noise_idx]] = advect_noise_vars[noise_idx]
set_dict[self._advect_noise_weight] = advect_noise_weight
if self._mask_noise_speed:
if len(noise_speed_masks) == 0:
for val in self._masked_noise_speed_weights.values():
set_dict[val] = np.zeros(val.shape)
else:
for key in noise_speed_masks.keys():
set_dict[self._masked_noise_speed_weights[
key]] = noise_speed_masks[key]
tflib.set_vars(set_dict)
else:
tflib.run(self._noise_init_op)
self._opt.reset_optimizer_state()
self._cur_step = 0
def training_loop(
submit_config,
G_args={}, # Options for generator network.
D_args={}, # Options for discriminator network.
G_opt_args={}, # Options for generator optimizer.
D_opt_args={}, # Options for discriminator optimizer.
G_loss_args={}, # Options for generator loss.
D_loss_args={}, # Options for discriminator loss.
dataset_args={}, # Options for dataset.load_dataset().
sched_args={}, # Options for train.TrainingSchedule.
grid_args={}, # Options for train.setup_snapshot_image_grid().
metric_arg_list=[], # Options for MetricGroup.
tf_config={}, # Options for tflib.init_tf().
G_smoothing_kimg=10.0, # Half-life of the running average of generator weights.
D_repeats=1, # How many times the discriminator is trained per G iteration.
minibatch_repeats=4, # Number of minibatches to run before adjusting training parameters.
reset_opt_for_new_lod=False, # Reset optimizer internal state (e.g. Adam moments) when new layers are introduced?
total_kimg=15000, # Total length of the training, measured in thousands of real images.
mirror_augment=False, # Enable mirror augment?
drange_net=[
-1, 1
], # Dynamic range used when feeding image data to the networks.
image_snapshot_ticks=1, # How often to export image snapshots?
network_snapshot_ticks=500, # How often to export network snapshots?
save_tf_graph=False, # Include full TensorFlow computation graph in the tfevents file?
save_weight_histograms=False, # Include weight histograms in the tfevents file?
resume_run_id=None, # Run ID or network pkl to resume training from, None = start from scratch.
resume_snapshot=None, # Snapshot index to resume training from, None = autodetect.
resume_kimg=0.0, # Assumed training progress at the beginning. Affects reporting and training schedule.
resume_time=0.0
): # Assumed wallclock time at the beginning. Affects reporting.
# Initialize dnnlib and TensorFlow.
ctx = dnnlib.RunContext(submit_config, train)
tflib.init_tf(tf_config)
# Load training set.
training_set = dataset.load_dataset(data_dir=config.data_dir,
verbose=True,
**dataset_args)
# Construct networks.
with tf.device('/gpu:0'):
if resume_run_id is not None:
network_pkl = misc.locate_network_pkl(resume_run_id,
resume_snapshot)
print('Loading networks from "%s"...' % network_pkl)
G, D, Gs = misc.load_pkl(network_pkl)
else:
print('Constructing networks...')
G = tflib.Network('G',
num_channels=training_set.shape[0],
resolution=training_set.shape[1],
label_size=training_set.label_size,
**G_args)
D = tflib.Network('D',
num_channels=training_set.shape[0],
resolution=training_set.shape[1],
label_size=training_set.label_size,
**D_args)
Gs = G.clone('Gs')
G.print_layers()
D.print_layers()
print('Building TensorFlow graph...')
with tf.name_scope('Inputs'), tf.device('/cpu:0'):
lod_in = tf.placeholder(tf.float32, name='lod_in', shape=[])
lrate_in = tf.placeholder(tf.float32, name='lrate_in', shape=[])
minibatch_in = tf.placeholder(tf.int32, name='minibatch_in', shape=[])
minibatch_split = minibatch_in // submit_config.num_gpus
Gs_beta = 0.5**tf.div(tf.cast(minibatch_in,
tf.float32), G_smoothing_kimg *
1000.0) if G_smoothing_kimg > 0.0 else 0.0
G_opt = tflib.Optimizer(name='TrainG',
learning_rate=lrate_in,
**G_opt_args)
D_opt = tflib.Optimizer(name='TrainD',
learning_rate=lrate_in,
**D_opt_args)
for gpu in range(submit_config.num_gpus):
with tf.name_scope('GPU%d' % gpu), tf.device('/gpu:%d' % gpu):
G_gpu = G if gpu == 0 else G.clone(G.name + '_shadow')
D_gpu = D if gpu == 0 else D.clone(D.name + '_shadow')
lod_assign_ops = [
tf.assign(G_gpu.find_var('lod'), lod_in),
tf.assign(D_gpu.find_var('lod'), lod_in)
]
reals, labels = training_set.get_minibatch_tf()
reals = process_reals(reals, lod_in, mirror_augment,
training_set.dynamic_range, drange_net)
with tf.name_scope('G_loss'), tf.control_dependencies(
lod_assign_ops):
G_loss = dnnlib.util.call_func_by_name(
G=G_gpu,
D=D_gpu,
opt=G_opt,
training_set=training_set,
minibatch_size=minibatch_split,
**G_loss_args)
with tf.name_scope('D_loss'), tf.control_dependencies(
lod_assign_ops):
D_loss = dnnlib.util.call_func_by_name(
G=G_gpu,
D=D_gpu,
opt=D_opt,
training_set=training_set,
minibatch_size=minibatch_split,
reals=reals,
labels=labels,
**D_loss_args)
G_opt.register_gradients(tf.reduce_mean(G_loss), G_gpu.trainables)
D_opt.register_gradients(tf.reduce_mean(D_loss), D_gpu.trainables)
G_train_op = G_opt.apply_updates()
D_train_op = D_opt.apply_updates()
Gs_update_op = Gs.setup_as_moving_average_of(G, beta=Gs_beta)
with tf.device('/gpu:0'):
try:
peak_gpu_mem_op = tf.contrib.memory_stats.MaxBytesInUse()
except tf.errors.NotFoundError:
peak_gpu_mem_op = tf.constant(0)
print('Setting up snapshot image grid...')
grid_size, grid_reals, grid_labels, grid_latents = misc.setup_snapshot_image_grid(
G, training_set, **grid_args)
sched = training_schedule(cur_nimg=total_kimg * 1000,
training_set=training_set,
num_gpus=submit_config.num_gpus,
**sched_args)
grid_fakes = Gs.run(grid_latents,
grid_labels,
is_validation=True,
minibatch_size=sched.minibatch //
submit_config.num_gpus)
print('Setting up run dir...')
misc.save_image_grid(grid_reals,
os.path.join(submit_config.run_dir, 'reals.png'),
drange=training_set.dynamic_range,
grid_size=grid_size)
misc.save_image_grid(grid_fakes,
os.path.join(submit_config.run_dir,
'fakes%06d.png' % resume_kimg),
drange=drange_net,
grid_size=grid_size)
summary_log = tf.summary.FileWriter(submit_config.run_dir)
if save_tf_graph:
summary_log.add_graph(tf.get_default_graph())
if save_weight_histograms:
G.setup_weight_histograms()
D.setup_weight_histograms()
metrics = metric_base.MetricGroup(metric_arg_list)
print('Training...\n')
ctx.update('', cur_epoch=resume_kimg, max_epoch=total_kimg)
maintenance_time = ctx.get_last_update_interval()
cur_nimg = int(resume_kimg * 1000)
cur_tick = 0
tick_start_nimg = cur_nimg
prev_lod = -1.0
while cur_nimg < total_kimg * 1000:
if ctx.should_stop(): break
# Choose training parameters and configure training ops.
sched = training_schedule(cur_nimg=cur_nimg,
training_set=training_set,
num_gpus=submit_config.num_gpus,
**sched_args)
training_set.configure(sched.minibatch // submit_config.num_gpus,
sched.lod)
if reset_opt_for_new_lod:
if np.floor(sched.lod) != np.floor(prev_lod) or np.ceil(
sched.lod) != np.ceil(prev_lod):
G_opt.reset_optimizer_state()
D_opt.reset_optimizer_state()
prev_lod = sched.lod
# Run training ops.
for _mb_repeat in range(minibatch_repeats):
for _D_repeat in range(D_repeats):
tflib.run(
[D_train_op, Gs_update_op], {
lod_in: sched.lod,
lrate_in: sched.D_lrate,
minibatch_in: sched.minibatch
})
cur_nimg += sched.minibatch
tflib.run(
[G_train_op], {
lod_in: sched.lod,
lrate_in: sched.G_lrate,
minibatch_in: sched.minibatch
})
# Perform maintenance tasks once per tick.
done = (cur_nimg >= total_kimg * 1000)
if cur_nimg >= tick_start_nimg + sched.tick_kimg * 1000 or done:
cur_tick += 1
tick_kimg = (cur_nimg - tick_start_nimg) / 1000.0
tick_start_nimg = cur_nimg
tick_time = ctx.get_time_since_last_update()
total_time = ctx.get_time_since_start() + resume_time
# Report progress.
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 %-4.1f'
% (autosummary('Progress/tick', cur_tick),
autosummary('Progress/kimg', cur_nimg / 1000.0),
autosummary('Progress/lod', sched.lod),
autosummary('Progress/minibatch', sched.minibatch),
dnnlib.util.format_time(
autosummary('Timing/total_sec', total_time)),
autosummary('Timing/sec_per_tick', tick_time),
autosummary('Timing/sec_per_kimg', tick_time / tick_kimg),
autosummary('Timing/maintenance_sec', maintenance_time),
autosummary('Resources/peak_gpu_mem_gb',
peak_gpu_mem_op.eval() / 2**30)))
autosummary('Timing/total_hours', total_time / (60.0 * 60.0))
autosummary('Timing/total_days', total_time / (24.0 * 60.0 * 60.0))
# Save snapshots.
if cur_tick % image_snapshot_ticks == 0 or done:
grid_fakes = Gs.run(grid_latents,
grid_labels,
is_validation=True,
minibatch_size=sched.minibatch //
submit_config.num_gpus)
misc.save_image_grid(grid_fakes,
os.path.join(
submit_config.run_dir,
'fakes%06d.png' % (cur_nimg // 1000)),
drange=drange_net,
grid_size=grid_size)
if cur_tick % network_snapshot_ticks == 0 or done or cur_tick == 1:
pkl = os.path.join(
submit_config.run_dir,
'network-snapshot-%06d.pkl' % (cur_nimg // 1000))
misc.save_pkl((G, D, Gs), pkl)
metrics.run(pkl,
run_dir=submit_config.run_dir,
num_gpus=submit_config.num_gpus,
tf_config=tf_config)
# Update summaries and RunContext.
metrics.update_autosummaries()
tflib.autosummary.save_summaries(summary_log, cur_nimg)
ctx.update('%.2f' % sched.lod,
cur_epoch=cur_nimg // 1000,
max_epoch=total_kimg)
maintenance_time = ctx.get_last_update_interval() - tick_time
# Write final results.
misc.save_pkl((G, D, Gs),
os.path.join(submit_config.run_dir, 'network-final.pkl'))
summary_log.close()
ctx.close()
def _evaluate(self, Gs, num_gpus):
minibatch_size = num_gpus * self.minibatch_per_gpu
# Construct TensorFlow graph for each GPU.
result_expr = []
for gpu_idx in range(num_gpus):
with tf.device('/gpu:%d' % gpu_idx):
Gs_clone = Gs.clone()
# Generate images.
latents = tf.random_normal([self.minibatch_per_gpu] +
Gs_clone.input_shape[1:])
dlatents = Gs_clone.components.mapping.get_output_for(
latents, None, is_validation=True)
images = Gs_clone.components.synthesis.get_output_for(
dlatents, is_validation=True, randomize_noise=True)
# Downsample to 256x256. The attribute classifiers were built for 256x256.
if images.shape[2] > 256:
factor = images.shape[2] // 256
images = tf.reshape(images, [
-1, images.shape[1], images.shape[2] // factor, factor,
images.shape[3] // factor, factor
])
images = tf.reduce_mean(images, axis=[3, 5])
# Run classifier for each attribute.
result_dict = dict(latents=latents, dlatents=dlatents[:, -1])
for attrib_idx in self.attrib_indices:
classifier = misc.load_pkl(classifier_urls[attrib_idx])
logits = classifier.get_output_for(images, None)
predictions = tf.nn.softmax(
tf.concat([logits, -logits], axis=1))
result_dict[attrib_idx] = predictions
result_expr.append(result_dict)
# Sampling loop.
results = []
for _ in range(0, self.num_samples, minibatch_size):
results += tflib.run(result_expr)
results = {
key: np.concatenate([value[key] for value in results], axis=0)
for key in results[0].keys()
}
# Calculate conditional entropy for each attribute.
conditional_entropies = defaultdict(list)
for attrib_idx in self.attrib_indices:
# Prune the least confident samples.
pruned_indices = list(range(self.num_samples))
pruned_indices = sorted(
pruned_indices, key=lambda i: -np.max(results[attrib_idx][i]))
pruned_indices = pruned_indices[:self.num_keep]
# Fit SVM to the remaining samples.
svm_targets = np.argmax(results[attrib_idx][pruned_indices],
axis=1)
for space in ['latents', 'dlatents']:
svm_inputs = results[space][pruned_indices]
try:
svm = sklearn.svm.LinearSVC()
svm.fit(svm_inputs, svm_targets)
svm.score(svm_inputs, svm_targets)
svm_outputs = svm.predict(svm_inputs)
except:
svm_outputs = svm_targets # assume perfect prediction
# Calculate conditional entropy.
p = [[
np.mean([
case == (row, col)
for case in zip(svm_outputs, svm_targets)
]) for col in (0, 1)
] for row in (0, 1)]
conditional_entropies[space].append(conditional_entropy(p))
