'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))