def write_block(self): """ Writes a block of images to disk :return: imgfn_2_imgid: dict mapping each image filename written to each image id written """ assert self._images is not None, 'no images loaded' if self._blocksize is not None: blocksize = self._blocksize else: blocksize = len(self._images) self._iblock += 1 self._iloop += 1 view = self._random_generator.permutation(self._nimgs) prioritized_view = np.argsort(self._remaining_times_toshow[view])[::-1][:blocksize] block_images = self._images[view[prioritized_view]] block_imgids = self._imgids[view[prioritized_view]] block_ids = ['block%03d_%03d' % (self._iblock, i) for i in range(blocksize)] block_imgfns = ['%s_%s.bmp' % (blockid, imgid) for blockid, imgid in zip(block_ids, block_imgids)] imgfn_2_imgid = {name: imgid for name, imgid in zip(block_imgfns, block_imgids)} # images written here utils.write_images(block_images, block_imgfns, self._writedir, self._imsize) self._curr_block_imgfns = block_imgfns self._imgfn_2_imgid = imgfn_2_imgid self._remaining_times_toshow[view[prioritized_view]] -= 1 return imgfn_2_imgid
def save_current_state(self, image_size=None): """ Saves current samples and images to disk """ self.save_current_codes() utils.write_images(self._curr_images, self._curr_sample_ids, self._logdir, image_size)
def upload(): num_uploads = get_num_uploads(request) if num_uploads == 0: return redirect(url_for('error', status='empty')) if num_uploads > MAX_UPLOADS: return redirect(url_for('error', status='limit_exceed')) if invalid_extension(request): return redirect(url_for('error', status='invalid_type')) session_id = uuid.uuid4().hex images = get_images(request) images = utils.change_ext(images, IMG_EXT) base_path = os.path.join(UPLOAD_FOLDER, session_id) prepared_img_path = os.path.join(base_path, 'prepared') visualised_img_path = get_visualised_imgs_dir(session_id) result_path = os.path.join(base_path) utils.write_images(images, prepared_img_path, ext=IMG_EXT) vesicle.detect(images, result_path, model) utils.write_images(images, visualised_img_path, ext=IMG_EXT) make_archive('result', 'zip', base_path) move('result.zip', base_path) return redirect(url_for('result', result_id=session_id))
def main(argv): del argv tf.gfile.MakeDirs(os.path.join(FLAGS.model_dir)) resolution = FLAGS.end_resolution initial_checkpoint = None while initial_checkpoint is None and resolution != 1: model_dir = os.path.join(FLAGS.model_dir, 'resolution_' + str(resolution)) initial_checkpoint = tf.train.latest_checkpoint(model_dir) resolution = resolution // 2 if initial_checkpoint is None or resolution == 1: resolution = FLAGS.start_resolution model_dir = os.path.join(FLAGS.model_dir, 'resolution_' + str(resolution)) else: resolution *= 2 model_dir = os.path.join(FLAGS.model_dir, 'resolution_' + str(resolution)) est, local_est = get_estimator(model_dir, resolution) current_step = estimator._load_global_step_from_checkpoint_dir(model_dir) # pylint: disable=protected-access,line-too-long tf.logging.info('Starting training for %d steps, current step: %d' % (FLAGS.train_steps, current_step)) while current_step < FLAGS.train_steps: if current_step != 0 and current_step % FLAGS.resolution_steps == 0 and resolution != FLAGS.end_resolution: resolution *= 2 tf.logging.info('Change of resolution from %d to %d' % (resolution // 2, resolution)) model_dir = os.path.join(FLAGS.model_dir, 'resolution_' + str(resolution)) change_resolution(resolution) est, local_est = get_estimator(model_dir, resolution) next_checkpoint = min(current_step + FLAGS.train_steps_per_eval, FLAGS.train_steps) est.train(input_fn=dataset.TrainInputFunction(FLAGS.noise_dim, resolution, 'NHWC'), max_steps=next_checkpoint) current_step = next_checkpoint tf.logging.info('Finished training step %d' % current_step) if FLAGS.eval_loss: metrics = est.evaluate( input_fn=dataset.TrainInputFunction(FLAGS.noise_dim, resolution, 'NHWC'), steps=FLAGS.num_eval_images // FLAGS.batch_size) tf.logging.info('Finished evaluating') tf.logging.info(metrics) generated_iter = local_est.predict(input_fn=noise_input_fn) images = [p['generated_images'][:, :, :] for p in generated_iter] filename = os.path.join( FLAGS.model_dir, '%s-%s.png' % (str(current_step).zfill(5), 'x' + str(resolution))) utils.write_images(images, filename, 'NHWC') tf.logging.info('Finished generating images')
def save_current_images(self, save_dir=None, **kwargs): """ Saves current images to disk """ if save_dir is None or not os.path.isdir(save_dir): save_dir = self._logdir utils.write_images(self._curr_images, self._curr_sample_ids, save_dir, **kwargs)
def write_block(self, wait_for_empty=-1): """ Writes a block of images to disk :param wait_for_empty: float if positive, check write_dir every so many seconds until it is empty :return: imgfn_2_imgid: dict mapping each image filename written to each image id written """ assert self._images is not None, 'no images loaded' blocksize = self._blocksize if self._blocksize is not None else len( self._images) self.iblock += 1 self._iloop += 1 try: wait_for_empty = float(wait_for_empty) if wait_for_empty > 0: while not self.check_write_dir_empty(): sleep(wait_for_empty) except ValueError: pass view = self._random_generator.permutation(self._nimgs) prioritized_view = np.argsort( self._remaining_times_toshow[view])[::-1][:blocksize] block_images = self._images[view[prioritized_view]] block_imgids = self._imgids[view[prioritized_view]] block_ids = [ f'block{self._iblock:03d}_{i:03d}' for i in range(blocksize) ] block_imgfns = [ f'{blockid}_{imgid}.{self._fmt}' for blockid, imgid in zip(block_ids, block_imgids) ] imgfn_2_imgid = { name: imgid for name, imgid in zip(block_imgfns, block_imgids) } # images written here utils.write_images(block_images, block_imgfns, self._writedir, size=self._imsize, grayscale=self._grayscale, fmt=self._fmt) self._curr_block_imgfns = block_imgfns self._imgfn_2_imgid = imgfn_2_imgid self._remaining_times_toshow[view[prioritized_view]] -= 1 return imgfn_2_imgid
def detect_from_dir(input_dir, output_dir, model: modellib.MaskRCNN): # Подготовка изображений prepared_images_dir = output_dir os.makedirs(prepared_images_dir, exist_ok=True) utils.prepare_images(input_dir, prepared_images_dir) images = utils.load_images(prepared_images_dir) begin = time.time() detect(images, output_dir, model) end = time.time() n = len(images) duration = end - begin print(f'Processed {n} in {duration} sec') print(f'{duration / n} sec per image') utils.write_images(images, os.path.join(output_dir, 'vis'))
def sample(model): samples, drgz, dcgz = model.sample(False) ### make session gpuConfig = tf.ConfigProto(gpu_options=tf.GPUOptions( allow_growth=True, per_process_gpu_memory_fraction=0.5)) sess = tf.Session(config=gpuConfig) sess.run(tf.global_variables_initializer()) saver = tf.train.Saver(tf.global_variables()) saver.restore(sess, "./checkpoints/model_{}".format(FLAGS.load_index)) DRGZ, DCGZ, s = sess.run([drgz, dcgz, samples]) write_images(s * 255., "./samples/manual") np.savetxt("./samples/drgz.txt", DRGZ) np.savetxt("./samples/dcgz.txt", DCGZ)
def save_test_images(self, test_display_data, curr_iter, is_test=True): st = time() overlay = self.overlay save_images = [] for i in range(len(test_display_data)): self.set_input_display(test_display_data[i]) self.test() save_images.append(self.image[0]) # Overlay mask: save_mask = (self.image[0] - (self.image[0] * self.mask[0].repeat(3, 1, 1)) + self.mask[0].repeat(3, 1, 1)) save_fake_mask = ( self.image[0] - (self.image[0] * self.fake_mask[0].repeat(3, 1, 1)) + self.fake_mask[0].repeat(3, 1, 1)) if overlay: save_images.append(save_mask) save_images.append(save_fake_mask) else: save_images.append(self.mask[0].repeat(3, 1, 1)) save_images.append(self.fake_mask[0].repeat(3, 1, 1)) for i in range(len(test_display_data)): # Append real masks (overlayed and itself): self.set_input_display(test_display_data[i]) self.test() save_images.append(self.r_im[0]) save_real_mask_seg = ( self.r_im[0] - (self.r_im[0] * self.r_mask[0].repeat(3, 1, 1)) + self.r_mask[0].repeat(3, 1, 1)) save_real_mask = ( self.r_im[0] - (self.r_im[0] * self.r_fake_mask[0].repeat(3, 1, 1)) + self.r_fake_mask[0].repeat(3, 1, 1)) save_images.append(save_real_mask_seg) save_images.append(save_real_mask) write_images( save_images, curr_iter, comet_exp=self.comet_exp, store_im=self.store_image, is_test=is_test, ) return time() - st
def validation_multi(model, criterion, valid_loader, num_classes, first_time, class_color_table, result_root): with torch.no_grad(): model.eval() losses = [] confusion_matrix = np.zeros((num_classes, num_classes), dtype=np.uint32) count = 0 for inputs, targets in valid_loader: inputs = utils.cuda(inputs) targets = utils.cuda(targets) outputs = model(inputs) loss = criterion(outputs, targets) losses.append(loss.item()) output_classes = outputs.data.cpu().numpy().argmax(axis=1) target_classes = targets.data.cpu().numpy() if first_time: utils.write_images(target_classes, class_color_table, result_root, "gt_", count) count = utils.write_images(output_classes, class_color_table, result_root, "pred_", count) confusion_matrix += calculate_confusion_matrix_from_arrays( output_classes, target_classes, num_classes) confusion_matrix = confusion_matrix[1:, 1:] # exclude background valid_loss = np.mean(losses) ious = { 'iou_{}'.format(cls + 1): iou for cls, iou in enumerate(calculate_iou(confusion_matrix)) } dices = { 'dice_{}'.format(cls + 1): dice for cls, dice in enumerate(calculate_dice(confusion_matrix)) } average_iou = np.mean(list(ious.values())) average_dices = np.mean(list(dices.values())) print('Valid loss: {:.4f}, average IoU: {:.4f}, average Dice: {:.4f}'. format(valid_loss, average_iou, average_dices)) metrics = {'valid_loss': valid_loss, 'iou': average_iou} metrics.update(ious) metrics.update(dices) return metrics
def generate_step(cfg, resolution): graph = tf.Graph() restore_dir = os.path.join(cfg.model_dir, 'resolution_' + str(resolution)) with graph.as_default(): # pylint: disable=E1129 input = input_pipelines.PredictInputFunction(cfg.noise_dim, resolution) params = {'data_dir': cfg.data_dir, 'batch_size': cfg.num_eval_images} features, labels = input(params) model = model_fn(features, labels, 'PREDICT', cfg) global_step = tf.train.get_or_create_global_step() with tf.Session() as sess: sess.run(tf.global_variables_initializer()) utils.restore(sess, restore_dir) images = sess.run(model) utils.write_images( images, cfg.model_dir + '/' + str(global_step.eval()).zfill(6) + '-' + str(resolution) + '.png', cfg.data_format) tf.reset_default_graph()
def save_test_images(self, test_display_data, curr_iter): save_images = [] for i in range(len(test_display_data)): self.set_input(Dict(test_display_data[i])) self.test() save_images.append(self.cloth[0]) save_images.append(self.mask[0].repeat(3, 1, 1)) save_images.append(self.image[0]) save_images.append(self.parse_cloth[0].repeat(3, 1, 1)) save_images.append(self.warped_cloth[0] * self.warped_mask[0]) save_images.append(self.warped_mask[0].repeat(3, 1, 1)) write_images( save_images, curr_iter, im_per_row=6, comet_exp=self.comet_exp, store_im=self.store_image, )
def write_block(self): assert self._images is not None, 'no images loaded' if self._blocksize is not None: blocksize = self._blocksize else: blocksize = len(self._images) self._iblock += 1 self._iloop += 1 view = self._random_generator.permutation(self._nimgs) prioritized_view = np.argsort(self._remaining_times_toshow[view])[::-1][:blocksize] block_images = self._images[view[prioritized_view]] block_imgids = self._imgids[view[prioritized_view]] block_ids = ['block%03d_%02d' % (self._iblock, i) for i in range(blocksize)] block_imgfns = ['%s_%s.bmp' % (blockid, imgid) for blockid, imgid in zip(block_ids, block_imgids)] imgfn_2_imgid = {name: imgid for name, imgid in zip(block_imgfns, block_imgids)} utils.write_images(block_images, block_imgfns, self._writedir, self._imsize) # the order of `block_ids` is not the same as `imgid` ! self._curr_block_imgfns = block_imgfns self._imgfn_2_imgid = imgfn_2_imgid self._remaining_times_toshow[view[prioritized_view]] -= 1 return imgfn_2_imgid
#manipulate first continious code manipulate_continious(idx, 0, decoded_samples, prior_samples, OUT_DIR + '/cvae_0_cont_code_manipul_gen.jpg') #manipulate second continious code manipulate_continious(idx, 1, decoded_samples, prior_samples, OUT_DIR + '/cvae_1_cont_code_manipul_gen.jpg') dataset_list = list(train_loader) images = dataset_list[0][0].cuda() labels = dataset_list[0][1].cuda() recons, enc_samples = trainer.vae.encode_decode(images, hard_catg=True) enc_samples_backup = enc_samples.clone().detach() #just save reconstructions and originals write_images([recons, images], 11, OUT_DIR, '_input_recons') idx = 1 #manipulate codes of encoded samples #manipulate first continious code manipulate_continious(idx, 0, recons, enc_samples, OUT_DIR + '/cvae_0_cont_code_manipul_rec.jpg') #manipulate second continious code manipulate_continious(idx, 1, recons, enc_samples, OUT_DIR + '/cvae_1_cont_code_manipul_rec.jpg') if opts.no_discrete: sys.exit('Done') #manipulate discrete code of encoded images and then decode out_list = []
if iblock != -1: t0 = time() # wait for matfile matfn = 'block%03d_code.mat' % iblock matfpath = os.path.join(respdir, matfn) print('waiting for %s' % matfn) while not os.path.isfile(matfpath): sleep(0.001) remove_bmp(respdir) copy_nature_image(natstimdir, expdir) sleep(0.9) # ensures mat file finish writing t1 = time() # load .mat file results imgid_list, codes = load_block_mat_code(matfpath) else: matfpath = initmat_path imgid_list, codes = load_block_mat_code(matfpath) imgid_list = ["gen_" + imgid for imgid in imgid_list ] # use gen as marker to distinguish from natural images copy_nature_image(natstimdir, expdir) # , prefix='block000') iblock += 1 # TODO Permutation and mixing can be added here ! But not needed! names = ['block%03d_' % (iblock) + id for id in imgid_list] # TODO More complex naming rule can be applied imgs = [generator.visualize(code) for code in codes] utils.write_images( imgs, names, respdir, format='bmp') # use jpg to compress size, use bmp to speed up utils.write_images(imgs, names, backupdir, format='jpg') copy(matfpath, backupdir)
real_annot).detach().cpu() # Evaluate generated images utils.show_images( fake_imgs, title=f"Batch of Fake Images After {epoch} Epochs") elif args.write_images is not None: # Write args.num_to_write images to files in ./generated/ device = torch.device(args.device) epoch, D, G, optimizerD, optimizerG = utils.load_checkpoint( f"checkpoints/checkpoint{args.write_images}.pt", device, ndf=ndf, ngf=ngf, latent_dims=latent_dims) dataloader = get_data_loader() with torch.no_grad(): images_generated = 0 while images_generated < args.num_to_write: for i, data in enumerate(dataloader): real_annots = data[1].type(torch.float).to( device) # 40 annotations for each image latent_vector = torch.randn( batch_size, latent_dims, 1, 1, device=device) # Generate noise fake_imgs = G(latent_vector, real_annots).detach().cpu( ) # Evaluate generated images utils.write_images(fake_imgs, images_generated) images_generated += batch_size if images_generated >= args.num_to_write: # Write num_to_write images in total break else: print(f"Command-line args: {args}") train(args) # Start training loop
def run(self): """ Main experiment loop """ self._load_nets() self._save_parameters() self.istep = 0 try: while self._max_steps is None or self.istep < self._max_steps: print(f'\n>>> step {self.istep:d}') last_codes = self.optimizer.current_samples_copy last_images = self.optimizer.current_images last_imgids = self.optimizer.current_image_ids last_scores = None t0 = time() if not self._cycle_refim_flags and self.istep == 0 and self.reference_images: # score images refscores = self.scorer.score(self.reference_images, self.reference_image_ids) t1 = time() # backup scores self.scorer.save_current_scores() t2 = time() # summarize scores & delays print( f'reference image scores: mean {np.nanmean(refscores)}, all {refscores}' ) print( f'step {self.istep:d} time: total {t2 - t0:.2f}s | wait for results {t1 - t0:.2f}s ' f'write records {t2 - t1:.2f}s') else: # score images syn_images = self.optimizer.current_images syn_image_ids = self.optimizer.current_image_ids if self.reference_images: combined_scores = self.scorer.score( syn_images + self.reference_images, syn_image_ids + self.reference_image_ids) refscores = combined_scores[self.optimizer.nsamples:] else: combined_scores = self.scorer.score( syn_images, syn_image_ids) synscores = combined_scores[:self.optimizer.nsamples] t1 = time() # before update, backup codes (optimizer) and scores (scorer) if self._write_codes: self.optimizer.save_current_codes() if hasattr(self.optimizer, 'save_current_genealogy'): self.optimizer.save_current_genealogy() last_scores = synscores self.scorer.save_current_scores() t2 = time() # use results to update optimizer self.optimizer.step(synscores) t3 = time() # summarize scores & delays print( f'synthetic img scores: mean {np.nanmean(synscores, axis=0)}, all {synscores}' ) if self.reference_images: print( f'reference image scores: mean {np.nanmean(refscores, axis=0)}, all {refscores}' ) print( f'step {self.istep:d} time: total {t3 - t0:.2f}s | ' + f'wait for results {t1 - t0:.2f}s write records {t2 - t1:.2f}s ' f'optimizer update {t3 - t2:.2f}s') # refresh reference images being shown self.check_refresh_ref_ims() self.logger.flush() self.istep += 1 if self._wait_each_step: sleep(self._wait_each_step) print('\nfinished <<<\n\n') # gracefully exit except KeyboardInterrupt: print() print('... keyboard interrupt') print(f'stopped at step {self.istep:d} <<<\n\n') # save final results when stopped try: ibest = None if last_scores is not None: ibest = np.argmax(last_scores) if not self._write_codes: if self._write_last_codes or (self._write_best_last_code and ibest is None): utils.write_codes(last_codes, last_imgids, self.logdir) elif self._write_best_last_code and ibest is not None: utils.write_codes([last_codes[ibest]], [last_imgids[ibest]], self.logdir) if self._write_last_images or (self._write_best_last_image and ibest is None): utils.write_images(last_images, last_imgids, self.logdir) elif self._write_best_last_image and ibest is not None: utils.write_images([last_images[ibest]], [last_imgids[ibest]], self.logdir) except NameError: pass self.logger.stop()
def main(): args = get_arguments() with open(args.model_params, 'r') as f: model_params = json.load(f) with open(args.training_params, 'r') as f: train_params = json.load(f) try: directories = validate_directories(args) except ValueError as e: print('Some arguments are wrong:') print(str(e)) return logdir = directories['logdir'] restore_from = directories['restore_from'] # Even if we restored the model, we will treat it as new training # if the trained model is written into an arbitrary location. is_overwritten_training = logdir != restore_from receptive_field = WaveNetModel.calculate_receptive_field( model_params['filter_width'], model_params['dilations'], model_params['initial_filter_width']) # Save arguments and model params into file save_run_config(args, receptive_field, STARTED_DATESTRING, logdir) # Create coordinator. coord = tf.train.Coordinator() # Create data loader. with tf.name_scope('create_inputs'): reader = WavMidReader(data_dir=args.data_dir_train, coord=coord, audio_sample_rate=model_params['audio_sr'], receptive_field=receptive_field, velocity=args.velocity, sample_size=args.sample_size, queues_size=(10, 10*args.batch_size)) data_batch = reader.dequeue(args.batch_size) # Create model. net = WaveNetModel( batch_size=args.batch_size, dilations=model_params['dilations'], filter_width=model_params['filter_width'], residual_channels=model_params['residual_channels'], dilation_channels=model_params['dilation_channels'], skip_channels=model_params['skip_channels'], output_channels=model_params['output_channels'], use_biases=model_params['use_biases'], initial_filter_width=model_params['initial_filter_width']) input_data = tf.placeholder(dtype=tf.float32, shape=(args.batch_size, None, 1)) input_labels = tf.placeholder(dtype=tf.float32, shape=(args.batch_size, None, model_params['output_channels'])) loss, probs = net.loss(input_data=input_data, input_labels=input_labels, pos_weight=train_params['pos_weight'], l2_reg_str=train_params['l2_reg_str']) optimizer = optimizer_factory[args.optimizer]( learning_rate=train_params['learning_rate'], momentum=train_params['momentum']) trainable = tf.trainable_variables() optim = optimizer.minimize(loss, var_list=trainable) # Set up logging for TensorBoard. writer = tf.summary.FileWriter(logdir) writer.add_graph(tf.get_default_graph()) run_metadata = tf.RunMetadata() summaries = tf.summary.merge_all() histograms = tf.summary.merge_all(key=HKEY) # Separate summary ops for validation, since they are # calculated only once per evaluation cycle. with tf.name_scope('validation_summaries'): metric_summaries = metrics_empty_dict() metric_value = tf.placeholder(tf.float32) for name in metric_summaries.keys(): metric_summaries[name] = tf.summary.scalar(name, metric_value) images_buffer = tf.placeholder(tf.string) images_batch = tf.stack( [tf.image.decode_png(images_buffer[0], channels=4), tf.image.decode_png(images_buffer[1], channels=4), tf.image.decode_png(images_buffer[2], channels=4)]) images_summary = tf.summary.image('estim', images_batch) audio_data = tf.placeholder(tf.float32) audio_summary = tf.summary.audio('input', audio_data, model_params['audio_sr']) # Set up session sess = tf.Session(config=tf.ConfigProto(log_device_placement=False)) init = tf.global_variables_initializer() sess.run(init) # Saver for storing checkpoints of the model. saver = tf.train.Saver(var_list=tf.trainable_variables(), max_to_keep=args.max_checkpoints) # Trainer for keeping best validation-performing model # and optional early stopping. trainer = Trainer(sess, logdir, train_params['early_stop_limit'], 0.999) try: saved_global_step = load(saver, sess, restore_from) if is_overwritten_training or saved_global_step is None: # The first training step will be saved_global_step + 1, # therefore we put -1 here for new or overwritten trainings. saved_global_step = -1 except: print('Something went wrong while restoring checkpoint. ' 'Training will be terminated to avoid accidentally ' 'overwriting the previous model.') raise threads = tf.train.start_queue_runners(sess=sess, coord=coord) reader.start_threads(sess) step = None last_saved_step = saved_global_step try: for step in range(saved_global_step + 1, train_params['num_steps']): waveform, pianoroll = sess.run([data_batch[0], data_batch[1]]) feed_dict = {input_data : waveform, input_labels : pianoroll} # Reload switches from file on each step with open(RUNTIME_SWITCHES, 'r') as f: switch = json.load(f) start_time = time.time() if switch['store_meta'] and step % switch['store_every'] == 0: # Slow run that stores extra information for debugging. print('Storing metadata') run_options = tf.RunOptions( trace_level=tf.RunOptions.FULL_TRACE) summary, loss_value, _ = sess.run( [summaries, loss, optim], feed_dict=feed_dict, options=run_options, run_metadata=run_metadata) writer.add_summary(summary, step) writer.add_run_metadata(run_metadata, 'step_{:04d}'.format(step)) tl = timeline.Timeline(run_metadata.step_stats) timeline_path = os.path.join(logdir, 'timeline.trace') with open(timeline_path, 'w') as f: f.write(tl.generate_chrome_trace_format(show_memory=True)) else: summary, loss_value, _ = sess.run([summaries, loss, optim], feed_dict=feed_dict) writer.add_summary(summary, step) duration = time.time() - start_time print('step {:d} - loss = {:.3f}, ({:.3f} sec/step)' .format(step, loss_value, duration)) if step % switch['checkpoint_every'] == 0: save(saver, sess, logdir, step) last_saved_step = step # Evaluate model performance on validation data if step % switch['evaluate_every'] == 0: if switch['histograms']: hist_summary = sess.run(histograms) writer.add_summary(hist_summary, step) print('evaluating...') stats = 0, 0, 0, 0, 0, 0 est = np.empty([0, model_params['output_channels']]) ref = np.empty([0, model_params['output_channels']]) b_data, b_labels, b_cntr = ( np.empty((0, args.sample_size + receptive_field - 1, 1)), np.empty((0, model_params['output_channels'])), args.batch_size) # if (batch_size * sample_size > valid_data) single_pass() again while est.size == 0: # and ref.size == 0 and sum(stats) == 0 ... for data, labels in reader.single_pass( sess, args.data_dir_valid): # cumulate batch if b_cntr > 1: b_data, b_labels, decr = cumulateBatch( data, labels, b_data, b_labels) b_cntr -= decr continue elif args.batch_size > 1: b_data, b_labels, decr = cumulateBatch( data, labels, b_data, b_labels) if not decr: continue data = b_data labels = b_labels # reset batch cumulation variables b_data, b_labels, b_cntr = ( np.empty(( 0, args.sample_size + receptive_field - 1, 1 )), np.empty((0, model_params['output_channels'])), args.batch_size) predictions = sess.run( probs, feed_dict={input_data : data}) # Aggregate sums for metrics calculation stats_chunk = calc_stats( predictions, labels, args.threshold) stats = tuple([sum(x) for x in zip(stats, stats_chunk)]) est = np.append(est, predictions, axis=0) ref = np.append(ref, labels, axis=0) metrics = calc_metrics(None, None, None, stats=stats) write_metrics(metrics, metric_summaries, metric_value, writer, step, sess) trainer.check(metrics['f1_measure']) # Render evaluation results if switch['log_image'] or switch['log_sound']: sub_fac = int(model_params['audio_sr']/switch['midi_sr']) est = roll_subsample(est.T, sub_fac) ref = roll_subsample(ref.T, sub_fac) if switch['log_image']: write_images(est, ref, switch['midi_sr'], args.threshold, (8, 6), images_summary, images_buffer, writer, step, sess) if switch['log_sound']: write_audio(est, ref, switch['midi_sr'], model_params['audio_sr'], 0.007, audio_summary, audio_data, writer, step, sess) except KeyboardInterrupt: # Introduce a line break after ^C is displayed so save message # is on its own line. print() finally: if step > last_saved_step: save(saver, sess, logdir, step) coord.request_stop() coord.join(threads) flush_n_close(writer, sess)
def main(argv): del argv tpu_cluster_resolver = None if FLAGS.use_tpu: tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver( FLAGS.tpu, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project) config = tpu_config.RunConfig( cluster=tpu_cluster_resolver, model_dir=FLAGS.model_dir, tpu_config=tpu_config.TPUConfig( num_shards=FLAGS.num_shards, iterations_per_loop=FLAGS.iterations_per_loop)) est = tpu_estimator.TPUEstimator(model_fn=model_fn, use_tpu=FLAGS.use_tpu, config=config, params={"data_dir": FLAGS.data_dir}, train_batch_size=FLAGS.batch_size, eval_batch_size=FLAGS.batch_size) local_est = tpu_estimator.TPUEstimator( model_fn=model_fn, use_tpu=False, config=config, params={"data_dir": FLAGS.data_dir}, predict_batch_size=FLAGS.num_eval_images) else: est = tf.estimator.Estimator(model_fn=model_fn, model_dir=FLAGS.model_dir, params={ "data_dir": FLAGS.data_dir, "batch_size": FLAGS.batch_size }) local_est = tf.estimator.Estimator(model_fn=model_fn, model_dir=FLAGS.model_dir, params={ "data_dir": FLAGS.data_dir, "batch_size": FLAGS.num_eval_images }) tf.gfile.MakeDirs(os.path.join(FLAGS.model_dir)) if FLAGS.mode == 'train': tf.gfile.MakeDirs(os.path.join(FLAGS.model_dir, 'generated_images')) else: tf.gfile.MakeDirs(os.path.join(FLAGS.model_dir, 'sampled_images')) tf.gfile.MakeDirs(os.path.join(FLAGS.model_dir, 'code')) for _file in glob.glob(r'*.py'): shutil.copy(_file, os.path.join(FLAGS.model_dir, 'code')) current_step = estimator._load_global_step_from_checkpoint_dir( FLAGS.model_dir) # pylint: disable=protected-access,line-too-long if FLAGS.mode == 'train': tf.logging.info('Starting training for %d steps, current step: %d' % (FLAGS.train_steps, current_step)) while current_step < FLAGS.train_steps: next_checkpoint = min(current_step + FLAGS.train_steps_per_eval, FLAGS.train_steps) est.train(input_fn=dataset.InputFunction(True, FLAGS.noise_dim), max_steps=next_checkpoint) current_step = next_checkpoint tf.logging.info('Finished training step %d' % current_step) generated_iter = local_est.predict(input_fn=noise_input_fn_fixed) images = [] for _ in range(FLAGS.num_eval_images): p = next(generated_iter) images.append(p['generated_images'][:, :, :]) filename = os.path.join( FLAGS.model_dir, 'generated_images', 'gen_%s.jpg' % (str(current_step).zfill(5))) utils.write_images(images, filename, [-1, 1], 'RGB', 'JPEG') tf.logging.info('Finished generating images') elif FLAGS.mode == 'pred': count = 0 while count < 10: generated_iter = local_est.predict(input_fn=noise_input_fn) images = [] for _ in range(FLAGS.num_eval_images): p = next(generated_iter) images.append(p['generated_images'][:, :, :]) filename = os.path.join( FLAGS.model_dir, 'sampled_images', '%s_%s.jpg' % (str(current_step).zfill(5), str(count))) utils.write_images(images, filename, [-1, 1], 'RGB', 'JPEG') count += 1
def save_test_images(self, test_display_data, curr_iter, is_test=True): st = time() overlay = self.overlay save_images = [] for i in range(len(test_display_data)): self.set_input_display(test_display_data[i]) self.test() save_images.append(self.image[0]) # Overlay mask: save_mask = (self.image[0] - (self.image[0] * self.mask[0].repeat(3, 1, 1)) + self.mask[0].repeat(3, 1, 1)) """ print("-------------------------") print("keys: ", test_display_data[i].data.keys()) print("x: ", test_display_data[i].data.x.shape) print("m: ", test_display_data[i].data.m.shape) print("d: ", test_display_data[i].data.d.shape) print("self.mask: ", self.mask.shape) print("self.fake_mask: ", self.fake_mask) print("-------------------------") """ save_fake_mask = ( self.image[0] - (self.image[0] * self.fake_mask[0].repeat(3, 1, 1)) + self.fake_mask[0].repeat(3, 1, 1)) if overlay: save_images.append(save_mask) save_images.append(save_fake_mask) else: save_images.append(self.mask[0].repeat(3, 1, 1)) save_images.append(self.fake_mask[0].repeat(3, 1, 1)) save_images.append(self.depth[0].repeat(3, 1, 1)) for i in range(len(test_display_data)): # Append real masks (overlayed and itself): self.set_input_display(test_display_data[i]) self.test() save_images.append(self.r_im[0]) save_real_mask_seg = ( self.r_im[0] - (self.r_im[0] * self.r_mask[0].repeat(3, 1, 1)) + self.r_mask[0].repeat(3, 1, 1)) save_real_mask = ( self.r_im[0] - (self.r_im[0] * self.r_fake_mask[0].repeat(3, 1, 1)) + self.r_fake_mask[0].repeat(3, 1, 1)) save_images.append(save_real_mask_seg) save_images.append(save_real_mask) save_images.append(self.r_depth[0].repeat(3, 1, 1)) write_images( save_images, curr_iter, comet_exp=self.comet_exp, store_im=self.store_image, is_test=is_test, im_per_row=4, ) return time() - st
def main(): args = get_arguments() with open(args.model_params, 'r') as f: model_params = json.load(f) with open(args.training_params, 'r') as f: train_params = json.load(f) try: directories = validate_directories(args) except ValueError as e: print('Some arguments are wrong:') print(str(e)) return logdir = directories['logdir'] restore_from = directories['restore_from'] # Even if we restored the model, we will treat it as new training # if the trained model is written into an arbitrary location. is_overwritten_training = logdir != restore_from receptive_field = WaveNetModel.calculate_receptive_field( model_params['filter_width'], model_params['dilations'], model_params['initial_filter_width']) # Save arguments and model params into file save_run_config(args, receptive_field, STARTED_DATESTRING, logdir) # Create coordinator. coord = tf.train.Coordinator() # Create data loader. with tf.name_scope('create_inputs'): reader = WavMidReader(data_dir=args.data_dir_train, coord=coord, audio_sample_rate=model_params['audio_sr'], receptive_field=receptive_field, velocity=args.velocity, sample_size=args.sample_size, queues_size=(10, 10 * args.batch_size)) data_batch = reader.dequeue(args.batch_size) # Create model. net = WaveNetModel( batch_size=args.batch_size, dilations=model_params['dilations'], filter_width=model_params['filter_width'], residual_channels=model_params['residual_channels'], dilation_channels=model_params['dilation_channels'], skip_channels=model_params['skip_channels'], output_channels=model_params['output_channels'], use_biases=model_params['use_biases'], initial_filter_width=model_params['initial_filter_width']) input_data = tf.placeholder(dtype=tf.float32, shape=(args.batch_size, None, 1)) input_labels = tf.placeholder(dtype=tf.float32, shape=(args.batch_size, None, model_params['output_channels'])) loss, probs = net.loss(input_data=input_data, input_labels=input_labels, pos_weight=train_params['pos_weight'], l2_reg_str=train_params['l2_reg_str']) optimizer = optimizer_factory[args.optimizer]( learning_rate=train_params['learning_rate'], momentum=train_params['momentum']) trainable = tf.trainable_variables() optim = optimizer.minimize(loss, var_list=trainable) # Set up logging for TensorBoard. writer = tf.summary.FileWriter(logdir) writer.add_graph(tf.get_default_graph()) run_metadata = tf.RunMetadata() summaries = tf.summary.merge_all() histograms = tf.summary.merge_all(key=HKEY) # Separate summary ops for validation, since they are # calculated only once per evaluation cycle. with tf.name_scope('validation_summaries'): metric_summaries = metrics_empty_dict() metric_value = tf.placeholder(tf.float32) for name in metric_summaries.keys(): metric_summaries[name] = tf.summary.scalar(name, metric_value) images_buffer = tf.placeholder(tf.string) images_batch = tf.stack([ tf.image.decode_png(images_buffer[0], channels=4), tf.image.decode_png(images_buffer[1], channels=4), tf.image.decode_png(images_buffer[2], channels=4) ]) images_summary = tf.summary.image('estim', images_batch) audio_data = tf.placeholder(tf.float32) audio_summary = tf.summary.audio('input', audio_data, model_params['audio_sr']) # Set up session sess = tf.Session(config=tf.ConfigProto(log_device_placement=False)) init = tf.global_variables_initializer() sess.run(init) # Saver for storing checkpoints of the model. saver = tf.train.Saver(var_list=tf.trainable_variables(), max_to_keep=args.max_checkpoints) # Trainer for keeping best validation-performing model # and optional early stopping. trainer = Trainer(sess, logdir, train_params['early_stop_limit'], 0.999) try: saved_global_step = load(saver, sess, restore_from) if is_overwritten_training or saved_global_step is None: # The first training step will be saved_global_step + 1, # therefore we put -1 here for new or overwritten trainings. saved_global_step = -1 except: print('Something went wrong while restoring checkpoint. ' 'Training will be terminated to avoid accidentally ' 'overwriting the previous model.') raise threads = tf.train.start_queue_runners(sess=sess, coord=coord) reader.start_threads(sess) step = None last_saved_step = saved_global_step try: for step in range(saved_global_step + 1, train_params['num_steps']): waveform, pianoroll = sess.run([data_batch[0], data_batch[1]]) feed_dict = {input_data: waveform, input_labels: pianoroll} # Reload switches from file on each step with open(RUNTIME_SWITCHES, 'r') as f: switch = json.load(f) start_time = time.time() if switch['store_meta'] and step % switch['store_every'] == 0: # Slow run that stores extra information for debugging. print('Storing metadata') run_options = tf.RunOptions( trace_level=tf.RunOptions.FULL_TRACE) summary, loss_value, _ = sess.run([summaries, loss, optim], feed_dict=feed_dict, options=run_options, run_metadata=run_metadata) writer.add_summary(summary, step) writer.add_run_metadata(run_metadata, 'step_{:04d}'.format(step)) tl = timeline.Timeline(run_metadata.step_stats) timeline_path = os.path.join(logdir, 'timeline.trace') with open(timeline_path, 'w') as f: f.write(tl.generate_chrome_trace_format(show_memory=True)) else: summary, loss_value, _ = sess.run([summaries, loss, optim], feed_dict=feed_dict) writer.add_summary(summary, step) duration = time.time() - start_time print('step {:d} - loss = {:.3f}, ({:.3f} sec/step)'.format( step, loss_value, duration)) if step % switch['checkpoint_every'] == 0: save(saver, sess, logdir, step) last_saved_step = step # Evaluate model performance on validation data if step % switch['evaluate_every'] == 0: if switch['histograms']: hist_summary = sess.run(histograms) writer.add_summary(hist_summary, step) print('evaluating...') stats = 0, 0, 0, 0, 0, 0 est = np.empty([0, model_params['output_channels']]) ref = np.empty([0, model_params['output_channels']]) b_data, b_labels, b_cntr = (np.empty( (0, args.sample_size + receptive_field - 1, 1)), np.empty((0, model_params['output_channels'])), args.batch_size) # if (batch_size * sample_size > valid_data) single_pass() again while est.size == 0: # and ref.size == 0 and sum(stats) == 0 ... for data, labels in reader.single_pass( sess, args.data_dir_valid): # cumulate batch if b_cntr > 1: b_data, b_labels, decr = cumulateBatch( data, labels, b_data, b_labels) b_cntr -= decr continue elif args.batch_size > 1: b_data, b_labels, decr = cumulateBatch( data, labels, b_data, b_labels) if not decr: continue data = b_data labels = b_labels # reset batch cumulation variables b_data, b_labels, b_cntr = ( np.empty( (0, args.sample_size + receptive_field - 1, 1)), np.empty((0, model_params['output_channels'])), args.batch_size) predictions = sess.run(probs, feed_dict={input_data: data}) # Aggregate sums for metrics calculation stats_chunk = calc_stats(predictions, labels, args.threshold) stats = tuple( [sum(x) for x in zip(stats, stats_chunk)]) est = np.append(est, predictions, axis=0) ref = np.append(ref, labels, axis=0) metrics = calc_metrics(None, None, None, stats=stats) write_metrics(metrics, metric_summaries, metric_value, writer, step, sess) trainer.check(metrics['f1_measure']) # Render evaluation results if switch['log_image'] or switch['log_sound']: sub_fac = int(model_params['audio_sr'] / switch['midi_sr']) est = roll_subsample(est.T, sub_fac) ref = roll_subsample(ref.T, sub_fac) if switch['log_image']: write_images(est, ref, switch['midi_sr'], args.threshold, (8, 6), images_summary, images_buffer, writer, step, sess) if switch['log_sound']: write_audio(est, ref, switch['midi_sr'], model_params['audio_sr'], 0.007, audio_summary, audio_data, writer, step, sess) except KeyboardInterrupt: # Introduce a line break after ^C is displayed so save message # is on its own line. print() finally: if step > last_saved_step: save(saver, sess, logdir, step) coord.request_stop() coord.join(threads) flush_n_close(writer, sess)
def train(model): is_train = tf.placeholder(tf.bool, shape=(), name="phase") n_class = FLAGS.n_class if FLAGS.from_example == "mnist": n_train = 60000 n_valid = 10000 feeder, x, c = data_from_example("mnist") else: n_train = 73288 n_valid = 8158 train_x, train_c = batch_reader(FLAGS.train_csv) if FLAGS.valid_csv: valid_x, valid_c = batch_reader(FLAGS.valid_csv) x, c = tf.cond(is_train, lambda: [train_x, train_c], lambda: [valid_x, valid_c]) else: x, c = train_x, train_c drx, drgz, dcx, dcgz, gz = model(x, is_train) true_positive = tf.reduce_mean(tf.round(drx)) * 100. true_negative = tf.reduce_mean(1 - tf.round(drgz)) * 100. class_accuracy = tf.reduce_mean( tf.cast(tf.equal(tf.argmax(dcx, axis=1), c), tf.float32)) * 100. entropy = tf.reduce_mean(tf.reduce_sum(-dcgz * tf.log(dcgz), axis=1)) with tf.name_scope('loss'): c_onehot = tf.one_hot(c, n_class) if FLAGS.wgan: d_loss = -tf.reduce_mean(drx) + tf.reduce_mean(drgz) g_loss = -tf.reduce_mean(drgz) else: d_loss = -tf.reduce_mean(tf.log(drx))\ -tf.reduce_mean(tf.log(1. - drgz)) if FLAGS.alternative_g_loss: g_loss = tf.reduce_mean(-tf.log(drgz)) else: g_loss = tf.reduce_mean(tf.log(1. - drgz)) if FLAGS.can: d_loss += -tf.reduce_mean( tf.reduce_sum(c_onehot * tf.log(dcx), axis=1)) g_loss += -tf.reduce_mean( tf.reduce_sum(1. / float(n_class) * tf.log(dcgz) + (1. - 1. / float(n_class)) * tf.log(1. - dcgz), axis=1)) update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) with tf.control_dependencies(update_ops): d_var_list = model.discriminator.vars + model.discriminator_head.vars g_var_list = model.generator.vars if FLAGS.wgan: d_opt = tf.train.RMSPropOptimizer( learning_rate=FLAGS.learning_rate) g_opt = tf.train.RMSPropOptimizer( learning_rate=FLAGS.learning_rate) else: d_opt = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate) g_opt = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate) d_gv = d_opt.compute_gradients(d_loss, d_var_list) d_step = d_opt.apply_gradients(d_gv) g_gv = g_opt.compute_gradients(g_loss, g_var_list) g_step = g_opt.apply_gradients(g_gv) if FLAGS.wgan: d_clip = [ v.assign(tf.clip_by_value(v, -0.01, 0.01)) for v in d_var_list ] tf.summary.image("samples", gz, max_outputs=20) tf.summary.scalar("true_positive", true_positive) tf.summary.scalar("true_negative", true_negative) tf.summary.scalar("class_accuracy", class_accuracy) tf.summary.scalar("entropy", entropy) tf.summary.scalar("discriminator_loss", d_loss) tf.summary.scalar("generator_loss", g_loss) w_list = [v for v in tf.global_variables() if "weights" in v.name] for w in w_list: tf.summary.histogram(w.name, w) b_list = [v for v in tf.global_variables() if "biases" in v.name] for b in b_list: tf.summary.histogram(b.name, b) gw_list = [gv[0] for gv in d_gv if "weights" in gv[1].name and gv[0] is not None] +\ [gv[0] for gv in g_gv if "weights" in gv[1].name and gv[0] is not None] for gw in gw_list: tf.summary.histogram(gw.name, gw) gb_list = [gv[0] for gv in d_gv if "biases" in gv[1].name and gv[0] is not None] +\ [gv[0] for gv in g_gv if "biases" in gv[1].name and gv[0] is not None] for gb in gb_list: tf.summary.histogram(gb.name, gb) summary = tf.summary.merge_all() ### make session gpuConfig = tf.ConfigProto(gpu_options=tf.GPUOptions( allow_growth=True, per_process_gpu_memory_fraction=0.5)) sess = tf.Session(config=gpuConfig) sess.run(tf.global_variables_initializer()) writer = tf.summary.FileWriter('summary', sess.graph) coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess=sess, coord=coord) saver = tf.train.Saver(tf.global_variables()) itersize = n_train / FLAGS.batch_size previous_time = time.time() if FLAGS.wgan: FLAGS.epoch *= 6 itersize /= 6 def make_train_feed_dict(phase): if FLAGS.from_example == "mnist": batch_x, batch_c = feeder.train.next_batch(FLAGS.batch_size) batch_x = np.reshape(batch_x, [-1, 28, 28, 1]) return {is_train: phase, x: batch_x, c: batch_c} else: return {is_train: phase} for step in range(FLAGS.epoch): error = 0. for itr in range(itersize): if FLAGS.wgan: for _ in range(5): feed_dict = make_train_feed_dict(True) _d, dl = sess.run([d_step, d_loss], feed_dict=feed_dict) sess.run(d_clip, feed_dict={is_train: True}) feed_dict = make_train_feed_dict(True) _g, gl = sess.run([g_step, g_loss], feed_dict=feed_dict) else: feed_dict = make_train_feed_dict(True) _d, _g, dl, gl = sess.run([d_step, g_step, d_loss, g_loss], feed_dict=feed_dict) current_time = time.time() #print "epoch:{}...{}/{} d_loss:{}, g_loss:{}, time:{}".format(step, itr, itersize, dl, gl, current_time-previous_time) previous_time = current_time if FLAGS.from_example == "mnist": batch_x, batch_c = feeder.test.next_batch(FLAGS.batch_size) batch_x = np.reshape(batch_x, [-1, 28, 28, 1]) feed_dict = {is_train: False, x: batch_x, c: batch_c} else: feed_dict = {is_train: False} tp, tn, ca, dl, gl, e, samples, s = sess.run([ true_positive, true_negative, class_accuracy, d_loss, g_loss, entropy, gz, summary ], feed_dict=feed_dict) writer.add_summary(s, step) current_time = time.time() in_pix_std = np.mean(np.std(samples, axis=0)) print "epoch:{}...validation d_loss:{} g_loss:{} true positive: {}%, true negative: {}%, class accuracy: {}%, entropy:{}, std_per_pixel: {} time:{}".format( step, dl, gl, tp, tn, ca, e, in_pix_std, current_time - previous_time) previous_time = current_time if (step + 1) % FLAGS.save_interval == 0: saver.save(sess, "{}_{}".format("checkpoints/model", step)) if (step + 1) % FLAGS.sample_interval == 0: write_images(samples * 255., "./samples/epoch_%.4d/" % step)
def main(argv): (opts, args) = parser.parse_args(argv) cudnn.benchmark = True model_name = os.path.splitext(os.path.basename(opts.config))[0] # Load experiment setting config = get_config(opts.config) max_iter = config['max_iter'] display_size = config['display_size'] # Setup model and data loader trainer = MUNIT_Trainer(config) trainer.cuda() train_loader_a, train_loader_b, test_loader_a, test_loader_b = get_all_data_loaders( config) test_display_images_a = Variable(torch.stack( [test_loader_a.dataset[i] for i in range(display_size)]).cuda(), volatile=True) test_display_images_b = Variable(torch.stack( [test_loader_b.dataset[i] for i in range(display_size)]).cuda(), volatile=True) train_display_images_a = Variable(torch.stack( [train_loader_a.dataset[i] for i in range(display_size)]).cuda(), volatile=True) train_display_images_b = Variable(torch.stack( [train_loader_b.dataset[i] for i in range(display_size)]).cuda(), volatile=True) # Setup logger and output folders train_writer = tensorboard.SummaryWriter(os.path.join( opts.log, model_name)) output_directory = os.path.join(opts.outputs, model_name) checkpoint_directory, image_directory = prepare_sub_folder( output_directory) shutil.copy(opts.config, os.path.join( output_directory, 'config.yaml')) # copy config file to output folder # Start training iterations = trainer.resume(checkpoint_directory) if opts.resume else 0 while True: for it, (images_a, images_b) in enumerate(izip(train_loader_a, train_loader_b)): trainer.update_learning_rate() images_a, images_b = Variable(images_a.cuda()), Variable( images_b.cuda()) # Main training code trainer.dis_update(images_a, images_b, config) trainer.gen_update(images_a, images_b, config) # Dump training stats in log file if (iterations + 1) % config['log_iter'] == 0: print("Iteration: %08d/%08d" % (iterations + 1, max_iter)) write_loss(iterations, trainer, train_writer) # Write images if (iterations + 1) % config['image_save_iter'] == 0: # Test set images image_outputs = trainer.sample(test_display_images_a, test_display_images_b) write_images( image_outputs, display_size, '%s/gen_test%08d.jpg' % (image_directory, iterations + 1)) # Train set images image_outputs = trainer.sample(train_display_images_a, train_display_images_b) write_images( image_outputs, display_size, '%s/gen_train%08d.jpg' % (image_directory, iterations + 1)) # HTML write_html(output_directory + "/index.html", iterations + 1, config['image_save_iter'], 'images') if (iterations + 1) % config['image_save_iter'] == 0: image_outputs = trainer.sample(test_display_images_a, test_display_images_b) write_images(image_outputs, display_size, '%s/gen.jpg' % image_directory) # Save network weights if (iterations + 1) % config['snapshot_save_iter'] == 0: trainer.save(checkpoint_directory, iterations) iterations += 1 if iterations >= max_iter: return
reconstructed = trainer.update_vae(images, config) if trainer.update_dis is not None: trainer.update_dis(images, config) torch.cuda.synchronize() # Dump training stats in log file if (iterations + 1) % config['log_iter'] == 0: print("Iteration: %08d/%08d" % (iterations + 1, max_iter)) write_loss(iterations, trainer, train_writer) #save some image stuff if (iterations + 1) % config['image_save_iter'] == 0 or iterations == 0: with torch.no_grad(): generated_images = trainer.vae.decoder(prior_samples) write_images([images, reconstructed, generated_images], display_size, image_directory, 'train_%08d' % (iterations + 1)) #visualize latent code influence on output if (iterations + 1) % (config['image_save_iter'] * 2) == 0: trainer.get_latent_visualization( image_directory, 'train_%08d' % (iterations + 1), images, prior_samples) iterations += 1 if iterations >= max_iter: if trainer.save is not None: trainer.save(checkpoint_directory, iterations) sys.exit('Finished training')
def run(self): """ Main experiment loop """ self._load_nets() self._save_parameters() self.istep = 0 try: while self._max_steps is None or self.istep < self._max_steps: print('\n>>> step %d' % self.istep) last_codes = self.optimizer.current_samples_copy last_images = self.optimizer.current_images last_imgids = self.optimizer.current_image_ids last_scores = None t0 = time() # if any natural images to show, show all at the first step if self.istep == 0 and self._n_natstim_to_show > 0: # score images natscores = self.scorer.score(self.natural_stimuli, self.natural_stimuli_ids) t1 = time() # backup scores self.scorer.save_current_scores() t2 = time() # summarize scores & delays print('natural image scores: mean {}, all {}'.format( np.nanmean(natscores), natscores)) print( 'step %d time: total %.2fs | wait for results %.2fs write records %.2fs' % (self.istep, t2 - t0, t1 - t0, t2 - t1)) else: # score images synscores = self.scorer.score( self.optimizer.current_images, self.optimizer.current_image_ids) t1 = time() # before update, backup codes and scores if self._write_codes: self.optimizer.save_current_codes() if hasattr(self.optimizer, 'save_current_genealogy'): self.optimizer.save_current_genealogy() last_scores = synscores self.scorer.save_current_scores() t2 = time() # use results to update optimizer self.optimizer.step(synscores) t3 = time() # summarize scores & delays print('synthetic img scores: mean {}, all {}'.format( np.nanmean(synscores), synscores)) print(( 'step %d time: total %.2fs | ' + 'wait for results %.2fs write records %.2fs optimizer update %.2fs' ) % (self.istep, t3 - t0, t1 - t0, t2 - t1, t3 - t2)) self.logger.flush() self.istep += 1 print('\nfinished <<<\n\n') # gracefully exit except KeyboardInterrupt: print() print('... keyboard interrupt') print('stopped at step %d <<<\n\n' % self.istep) # save final results when stopped try: ibest = None if last_scores is not None: ibest = np.argmax(last_scores) if not self._write_codes: if self._write_last_codes or (self._write_best_last_code and ibest is None): utils.write_codes(last_codes, last_imgids, self.logdir) elif self._write_best_last_code and ibest is not None: utils.write_codes([last_codes[ibest]], [last_imgids[ibest]], self.logdir) if self._write_last_images or (self._write_best_last_image and ibest is None): utils.write_images(last_images, last_imgids, self.logdir) elif self._write_best_last_image and ibest is not None: utils.write_images([last_images[ibest]], [last_imgids[ibest]], self.logdir) except NameError: pass self.logger.stop()
def main(): args = get_arguments() if (args.logdir is not None and os.path.isdir(args.logdir)): logdir = args.logdir else: print('Argument --logdir=\'{}\' is not (but should be) ' 'a path to valid directory.'.format(args.logdir)) return with open(args.model_params, 'r') as f: model_params = json.load(f) with open(RUNTIME_SWITCHES, 'r') as f: switch = json.load(f) receptive_field = WaveNetModel.calculate_receptive_field( model_params['filter_width'], model_params['dilations'], model_params['initial_filter_width']) # Create coordinator. coord = tf.train.Coordinator() # Create data loader. with tf.name_scope('create_inputs'): reader = WavMidReader(data_dir=args.data_dir_test, coord=coord, audio_sample_rate=model_params['audio_sr'], receptive_field=receptive_field, velocity=args.velocity, sample_size=args.sample_size, queues_size=(100, 100*BATCH_SIZE)) # Create model. net = WaveNetModel( batch_size=BATCH_SIZE, dilations=model_params['dilations'], filter_width=model_params['filter_width'], residual_channels=model_params['residual_channels'], dilation_channels=model_params['dilation_channels'], skip_channels=model_params['skip_channels'], output_channels=model_params['output_channels'], use_biases=model_params['use_biases'], initial_filter_width=model_params['initial_filter_width']) input_data = tf.placeholder(dtype=tf.float32, shape=(BATCH_SIZE, None, 1)) input_labels = tf.placeholder(dtype=tf.float32, shape=(BATCH_SIZE, None, model_params['output_channels'])) _, probs = net.loss(input_data=input_data, input_labels=input_labels, pos_weight=1.0, l2_reg_str=None) # Set up session sess = tf.Session(config=tf.ConfigProto(log_device_placement=False)) init = tf.global_variables_initializer() sess.run(init) # Saver for storing checkpoints of the model. saver = tf.train.Saver(var_list=tf.trainable_variables()) try: load(saver, sess, logdir) except: print('Something went wrong while restoring checkpoint.') raise try: stats = 0, 0, 0, 0, 0, 0 est = np.empty([model_params['output_channels'], 0]) ref = np.empty([model_params['output_channels'], 0]) sub_fac = int(model_params['audio_sr']/switch['midi_sr']) for data, labels in reader.single_pass(sess, args.data_dir_test): predictions = sess.run(probs, feed_dict={input_data : data}) # Aggregate sums for metrics calculation stats_chunk = calc_stats(predictions, labels, args.threshold) stats = tuple([sum(x) for x in zip(stats, stats_chunk)]) est = np.append(est, roll_subsample(predictions.T, sub_fac), axis=1) ref = np.append(ref, roll_subsample(labels.T, sub_fac, b=True), axis=1) metrics = calc_metrics(None, None, None, stats=stats) write_metrics(metrics, None, None, None, None, None, logdir=logdir) # Save subsampled data for further arbitrary evaluation np.save(logdir+'/est.npy', est) np.save(logdir+'/ref.npy', ref) # Render evaluation results figsize=(int(args.plot_scale*est.shape[1]/switch['midi_sr']), int(args.plot_scale*model_params['output_channels']/12)) if args.media: write_images(est, ref, switch['midi_sr'], args.threshold, figsize, None, None, None, 0, None, noterange=(21, 109), legend=args.plot_legend, logdir=logdir) write_audio(est, ref, switch['midi_sr'], model_params['audio_sr'], 0.007, None, None, None, 0, None, logdir=logdir) except KeyboardInterrupt: # Introduce a line break after ^C is displayed so save message # is on its own line. print() finally: coord.request_stop()
def save_current_average_sample(self, save_dir=None, **kwargs): if save_dir is None or not os.path.isdir(save_dir): save_dir = self._logdir utils.write_images([self.visualize(self.current_average_sample)], [f'{self._thread_prefix}step{self.istep:03d}_mean'], save_dir, **kwargs)
print('Loading model from {}'.format(checkpoint_directory)) trainer.resume(checkpoint_directory, hyperparameters=config) prior_samples = trainer.vae.prior.sample_prior(config['batch_size']) trainer.vae.train() new_sample = prior_samples.clone().detach() sample_list = [] with torch.no_grad(): for i in range(11): new_sample.data[0:1, 0] = -5. + i out = trainer.vae.decoder(new_sample) sample_list.append(out[0:1]) all_samples = torch.cat(sample_list, 0) write_images([all_samples], 11, '.', '_test') ################################ #generate grid from latent space ################################ idx = 11 rng = 10 start = 3 step = start * 2. / rng sample_list = [] target_sample = prior_samples.clone().detach() with torch.no_grad(): for i in range(rng):
colors = colors.to(device) flows = flows.to(device) # cpu_flows = flows.data.cpu().numpy() # images = colors.data.cpu().numpy() # cpu_flows = np.moveaxis(cpu_flows, [0, 1, 2, 3], [0, 3, 1, 2]) # images = np.moveaxis(images, [0, 1, 2, 3], [0, 3, 1, 2]) # # cv2.imshow('flow HSV', utils.draw_hsv(cpu_flows[0])) # cv2.imshow("color", images[0] * 0.5 + 0.5) # cv2.waitKey() embeddings = model(colors) images = utils.draw_embeddings(colors, embeddings, 10) utils.write_images(images, root=results_root, file_prefix="embedding_epoch_" + str(epoch) + "_" + str(i) + "_") cv2.waitKey() loss = cross_pixel_loss(embeddings, flows) optimizer.zero_grad() loss.backward() optimizer.step() step += 1 tq.update(batch_size) losses.append(loss.item()) mean_loss = np.mean(losses[-report_each:]) tq.set_postfix(loss='{:.5f}'.format(mean_loss)) tq.set_postfix(loss='{:.5f}'.format(np.mean(losses))) tq.close() save(epoch + 1)