def estimator(A_val, y_batch_val, hparams): """Function that returns the estimated image""" best_keeper = utils.BestKeeper(hparams) feed_dict = {A: A_val, y_batch: y_batch_val} for i in range(hparams.num_random_restarts): sess.run([z_batch.initializer]) for j in range(hparams.max_update_iter): _, lr_val, total_loss_val, \ m_loss1_val, \ m_loss2_val, \ zp_loss_val = sess.run([update_op, learning_rate, total_loss, m_loss1, m_loss2, zp_loss], feed_dict=feed_dict) logging_format = 'rr {} iter {} lr {} total_loss {} m_loss1 {} m_loss2 {} zp_loss {}' print logging_format.format(i, j, lr_val, total_loss_val, m_loss1_val, m_loss2_val, zp_loss_val) if hparams.gif and ((j % hparams.gif_iter) == 0): images = sess.run(x_hat_batch, feed_dict=feed_dict) for im_num, image in enumerate(images): save_dir = '{0}/{1}/'.format(hparams.gif_dir, im_num) utils.set_up_dir(save_dir) save_path = save_dir + '{0}.png'.format(j) image = image.reshape(hparams.image_shape) save_image(image, save_path) x_hat_batch_val, total_loss_batch_val = sess.run( [x_hat_batch, total_loss_batch], feed_dict=feed_dict) best_keeper.report(x_hat_batch_val, total_loss_batch_val) return best_keeper.get_best()
def store_layerwise_activations(net_prototxt, model, phase, keys, n, dst_fpath): ''' phase = caffe.TRAIN proto = '/mnt/antares_raid/home/oliver/Scripts/autoencoder_v2/autoencoder_with_MLP/autoencoder_with_MLP_net.prototxt' model = '/mnt/antares_raid/home/oliver/Scripts/autoencoder_v2/autoencoder_with_MLP/snapshots_with_MLP/_iter_780000.caffemodel' lmdb_path = '/mnt/antares_raid/home/oliver/Scripts/autoencoder_v2/MNIST_lmdb/MNIST_TRAIN_60000_rot_lmdb/shuffled/' dst_fpath = "/mnt/antares_raid/home/oliver/Scripts/autoencoder_v2/autoencoder_with_MLP/results/res_train.hdf5" n = 780 ''' set_up_dir("/".join(dst_fpath.split("/")[:-1])+"/") print("path created: {}".format("/".join(dst_fpath.split("/")[:-1])+"/")) print("Destination: {}".format(dst_fpath)) print(net_prototxt, model, phase) net = caffe.Net(net_prototxt, model, phase) infer_to_h5_fixed_dims(net, keys, n, dst_fpath, preserve_batch=False) print('Done creating %s'%dst_fpath)
def estimator(A_val, y_batch_val, hparams): """Function that returns the estimated image""" best_keeper = utils.BestKeeper(hparams) if hparams.measurement_type == 'project': # if y_batch_val.shape[0]!=hparams.batch_size: # y_batch_val_tmp = np.zeros((hparams.batch_size,hparams.num_measurements)) # y_batch_val_tmp[:y_batch_val.shape[0],:] = y_batch_val # y_batch_val = y_batch_val_tmp # print('Smaller INPUT NUMBER')#Or change hparams on the fly feed_dict = {y_batch: y_batch_val} else: feed_dict = {A: A_val, y_batch: y_batch_val} for i in range(hparams.num_random_restarts): sess.run(opt_reinit_op) for j in range(hparams.max_update_iter): _, lr_val, total_loss_val, \ m_loss1_val, \ m_loss2_val, \ zp_loss_val = sess.run([update_op, learning_rate, total_loss, m_loss1, m_loss2, zp_loss], feed_dict=feed_dict) logging_format = 'rr {} iter {} lr {} total_loss {} m_loss1 {} m_loss2 {} zp_loss {}' print( logging_format.format(i, j, lr_val, total_loss_val, m_loss1_val, m_loss2_val, zp_loss_val)) #print('n_z is {}'.format(hparams.n_z)) if total_loss_val == m_loss2_val and zp_loss_val > 0 and hparams.zprior_weight > 0: raise ValueError('NONONO') if hparams.gif and ((j % hparams.gif_iter) == 0): images = sess.run(x_hat_batch, feed_dict=feed_dict) for im_num, image in enumerate(images): save_dir = '{0}/{1}/'.format(hparams.gif_dir, im_num) utils.set_up_dir(save_dir) save_path = save_dir + '{0}.png'.format(j) image = image.reshape(hparams.image_shape) save_image(image, save_path) x_hat_batch_val, total_loss_batch_val = sess.run( [x_hat_batch, total_loss_batch], feed_dict=feed_dict) best_keeper.report(x_hat_batch_val, total_loss_batch_val) return best_keeper.get_best()
def estimator(Tx_val, Rx_val, Pilot_val, hparams): """Function that returns the estimated image""" best_keeper = utils.BestKeeper(hparams) if hparams.measurement_type == 'project': feed_dict = {y_batch: y_batch_val} else: feed_dict = {Tx: Tx_val, Rx: Rx_val, Pilot: Pilot_val} for i in range(hparams.num_random_restarts): sess.run(opt_reinit_op) for j in range(hparams.max_update_iter): if hparams.gif and ((j % hparams.gif_iter) == 0): images = sess.run(x_hat_batch, feed_dict=feed_dict) for im_num, image in enumerate(images): save_dir = '{0}/{1}/'.format(hparams.gif_dir, im_num) utils.set_up_dir(save_dir) save_path = save_dir + '{0}.png'.format(j) image = image.reshape(hparams.image_shape) save_image(image, save_path) _, lr_val, total_loss_val, \ m_loss1_val, \ m_loss2_val, \ zp_loss_val = sess.run([update_op, learning_rate, total_loss, m_loss1, m_loss2, zp_loss], feed_dict=feed_dict) logging_format = 'rr {} iter {} lr {} total_loss {} m_loss1 {} m_loss2 {} zp_loss {}' print logging_format.format(i, j, lr_val, total_loss_val, m_loss1_val, m_loss2_val, zp_loss_val) H_hat_val, total_loss_val = sess.run([H_hat, total_loss], feed_dict=feed_dict) best_keeper.report(H_hat_val, total_loss_val) return best_keeper.get_best()
def main(hparams): # Set up some stuff according to hparams utils.set_up_dir(hparams.ckpt_dir) utils.set_up_dir(hparams.sample_dir) utils.print_hparams(hparams) # encode x_ph = tf.placeholder(tf.float32, [None, hparams.n_input], name='x_ph') z_mean, z_log_sigma_sq = model_def.encoder(hparams, x_ph, 'enc', reuse=False) # sample eps = tf.random_normal((hparams.batch_size, hparams.n_z), 0, 1, dtype=tf.float32) z_sigma = tf.sqrt(tf.exp(z_log_sigma_sq)) z = z_mean + z_sigma * eps # reconstruct logits, x_reconstr_mean = model_def.generator(hparams, z, 'gen', reuse=False) # generator sampler z_ph = tf.placeholder(tf.float32, [None, hparams.n_z], name='x_ph') _, x_sample = model_def.generator(hparams, z_ph, 'gen', reuse=True) # define loss and update op total_loss = model_def.get_loss(x_ph, logits, z_mean, z_log_sigma_sq) opt = tf.train.AdamOptimizer(learning_rate=hparams.learning_rate) update_op = opt.minimize(total_loss) # Sanity checks for var in tf.global_variables(): print var.op.name print '' # Get a new session sess = tf.Session() # Model checkpointing setup model_saver = tf.train.Saver() init_op = tf.global_variables_initializer() sess.run(init_op) # Attempt to restore variables from checkpoint start_epoch = utils.try_restore(hparams, sess, model_saver) # Get data iterator iterator = data_input.channel_data_iteratior(hparams) next_element = iterator.get_next() # Training for epoch in range(start_epoch + 1, hparams.training_epochs): avg_loss = 0.0 num_batches = hparams.num_samples // hparams.batch_size batch_num = 0 for i in range(num_batches): try: x_batch_val = sess.run(next_element['H_data']) x_batch_val = np.squeeze(x_batch_val) batch_num += 1 feed_dict = {x_ph: x_batch_val} _, loss_val = sess.run([update_op, total_loss], feed_dict=feed_dict) #print(loss_val) avg_loss += loss_val / hparams.num_samples * hparams.batch_size except tf.errors.OutOfRangeError: print("End of dataset") break if epoch % hparams.summary_epoch == 0: print "Epoch:", '%04d' % (epoch), 'Avg loss = {:.9f}'.format( avg_loss) if epoch % hparams.ckpt_epoch == 0: save_path = os.path.join(hparams.ckpt_dir, 'channel_vae_model') model_saver.save(sess, save_path, global_step=epoch) save_path = os.path.join(hparams.ckpt_dir, 'channel_vae_model') model_saver.save(sess, save_path, global_step=hparams.training_epochs - 1)
def estimator(A_val, y_val, hparams): """Function that returns the estimated image""" A = torch.Tensor(A_val).to(device) y = torch.Tensor(y_val).to(device) shuffled_idx = torch.randperm(hparams.num_measurements) best_keeper = utils.BestKeeper((hparams.batch_size, hparams.n_input)) best_keeper_z = utils.BestKeeper((hparams.batch_size, 512)) def sample(z): return model.test(z, getAvG=True, toCPU=(not hparams.cuda)) def get_loss(xf, xg): # compute measurements yf_batch = torch.mm(xf.view(batch_size, -1), A) yg_batch = torch.mm(xg.view(batch_size, -1), A) # compute corresponding losses loss_1 = se(yf_batch, y) loss_2 = se(yg_batch, y) loss_3 = loss_1 - loss_2 # now find median block of loss_1 - loss_2 loss_3 = loss_1 - loss_2 #shuffle the losses loss_3 = loss_3[:, shuffled_idx] loss_3 = loss_3[:, :mom_batch_size * ( A.shape[0] // mom_batch_size )] # make the number of rows a multiple of batch size loss_3 = loss_3.view(batch_size, -1, mom_batch_size) # reshape loss_3 = loss_3.mean(axis=-1) # find mean on each batch loss_3_numpy = loss_3.detach().cpu().numpy() # convert to numpy median_idx = np.argsort(loss_3_numpy, axis=1)[:, loss_3_numpy.shape[1] // 2] # sort and pick middle element # pick median block loss_batch = loss_3[range( batch_size), median_idx] # torch.mean(loss_1_mom - loss_2_mom) return loss_batch for i in range(hparams.num_random_restarts): zf_batch = torch.randn(hparams.batch_size, 512).to(device) zg_batch = torch.randn(hparams.batch_size, 512).to(device) z_output_batch = torch.zeros(hparams.batch_size, 512).to(device) zf_batch.requires_grad_() zg_batch.requires_grad_() opt1 = utils.get_optimizer(zf_batch, hparams.learning_rate, hparams) opt2 = utils.get_optimizer(zg_batch, hparams.learning_rate, hparams) for j in range(hparams.max_update_iter): xf_batch, xg_batch = sample(zf_batch), sample(zg_batch) if hparams.gif and ((j % hparams.gif_iter) == 0): images = xf_batch.detach().cpu().numpy() for im_num, image in enumerate(images): save_dir = '{0}/{1}/'.format(hparams.gif_dir, im_num) utils.set_up_dir(save_dir) save_path = save_dir + '{0}.png'.format(j) image = image.reshape(hparams.image_shape) save_image(image, save_path) opt1.zero_grad() xf_batch, xg_batch = sample(zf_batch), sample(zg_batch) loss_f_batch = get_loss(xf_batch, xg_batch) loss_f = loss_f_batch.mean() loss_f.backward() opt1.step() opt2.zero_grad() xf_batch, xg_batch = sample(zf_batch), sample(zg_batch) loss_g_batch = -1 * get_loss(xf_batch, xg_batch) loss_g = loss_g_batch.mean() loss_g.backward() opt2.step() logging_format = 'rr {} iter {} loss_f {} loss_g {}' print(logging_format.format(i, j, loss_f.item(), loss_g.item())) if j >= hparams.max_update_iter - 200: z_output_batch += zf_batch.detach() z_output_batch = z_output_batch / 200 x_hat_batch = sample(z_output_batch) y_hat_batch = torch.mm(x_hat_batch.view(hparams.batch_size, -1), A) m_loss_batch = get_loss(x_hat_batch, xg_batch) best_keeper.report( x_hat_batch.view(hparams.batch_size, -1).detach().cpu().numpy(), m_loss_batch.detach().cpu().numpy()) best_keeper_z.report( z_output_batch.view(hparams.batch_size, -1).detach().cpu().numpy(), m_loss_batch.detach().cpu().numpy()) return best_keeper.get_best(), best_keeper_z.get_best( ), best_keeper.losses_val_best
parser.add_argument('--lw_perc', type=float, default=1) # --- parser.add_argument('--save_interval', type=int, default=100) parser.add_argument('--print_interval', type=int, default=10) parser.add_argument('--epoch', type=int, default=20) parser.add_argument('-p', '--project_name', type=str, default="") parser.add_argument('--speedup', type=int, default=16) parser.add_argument('--debug', action="store_true", help="if debug, log will be printed to screen rather than saved") parser.add_argument('--screen', action="store_true", help="if print log to screen") parser.add_argument('--updim_relu', action="store_true", help="if use relu for the 1x1 conv") parser.add_argument('--mode', type=str, choices=['wct_se', 'wct_sd']) parser.add_argument('--stage', type=int, choices=[1,2,3,4,5]) args = parser.parse_args() # set up log dirs TimeID, ExpID, rec_img_path, weights_path, log = set_up_dir(args.project_name, args.resume, args.debug) logprint = LogPrint(log, ExpID, args.screen) args.ExpID = ExpID args.CodeID = get_CodeID() loghub = LogHub() # Set up model, data, optimizer if args.mode == "wct_se": args.BE = "trained_models/original_wct_models/vgg_normalised_conv%d_1.t7" % args.stage args.BD = "trained_models/our_BD/%dBD_E30S0.pth" % args.stage if args.pretrained_init: args.SE = "trained_models/small16x_ae_base/e%d_base.pth" % args.stage net = TrainSE_With_WCTDecoder(args).cuda() dataset = Dataset(args.content_train, args.shorter_side) train_loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=args.batch_size, shuffle=True)
def main(hparams): # Set up some stuff according to hparams utils.set_up_dir(hparams.ckpt_dir) utils.set_up_dir(hparams.sample_dir) utils.print_hparams(hparams) # encode x_ph = tf.placeholder(tf.float32, [None, hparams.n_input], name='x_ph') z_mean, z_log_sigma_sq = model_def.encoder(hparams, x_ph, 'enc', reuse=False) # sample eps = tf.random_normal((hparams.batch_size, hparams.n_z), 0, 1, dtype=tf.float32) z_sigma = tf.sqrt(tf.exp(z_log_sigma_sq)) z = z_mean + z_sigma * eps # reconstruct logits, x_reconstr_mean = model_def.generator(hparams, z, 'gen', reuse=False) # generator sampler z_ph = tf.placeholder(tf.float32, [None, hparams.n_z], name='x_ph') _, x_sample = model_def.generator(hparams, z_ph, 'gen', reuse=True) # define loss and update op total_loss = model_def.get_loss(x_ph, logits, z_mean, z_log_sigma_sq) opt = tf.train.AdamOptimizer(learning_rate=hparams.learning_rate) update_op = opt.minimize(total_loss) # Sanity checks for var in tf.global_variables(): print var.op.name print '' # Get a new session sess = tf.Session() # Model checkpointing setup model_saver = tf.train.Saver() init_op = tf.global_variables_initializer() sess.run(init_op) # Attempt to restore variables from checkpoint start_epoch = utils.try_restore(hparams, sess, model_saver) # Get data iterator iterator = data_input.omniglot_data_iterator() # Training for epoch in range(start_epoch + 1, hparams.training_epochs): avg_loss = 0.0 num_batches = hparams.num_samples // hparams.batch_size batch_num = 0 for (x_batch_val, _) in iterator(hparams, num_batches): batch_num += 1 feed_dict = {x_ph: x_batch_val} _, loss_val = sess.run([update_op, total_loss], feed_dict=feed_dict) avg_loss += loss_val / hparams.num_samples * hparams.batch_size if batch_num % 100 == 0: x_reconstr_mean_val = sess.run(x_reconstr_mean, feed_dict={x_ph: x_batch_val}) z_val = np.random.randn(hparams.batch_size, hparams.n_z) x_sample_val = sess.run(x_sample, feed_dict={z_ph: z_val}) utils.save_images( np.reshape(x_reconstr_mean_val, [-1, 28, 28]), [10, 10], '{}/reconstr_{:02d}_{:04d}.png'.format( hparams.sample_dir, epoch, batch_num)) utils.save_images( np.reshape(x_batch_val, [-1, 28, 28]), [10, 10], '{}/orig_{:02d}_{:04d}.png'.format(hparams.sample_dir, epoch, batch_num)) utils.save_images( np.reshape(x_sample_val, [-1, 28, 28]), [10, 10], '{}/sampled_{:02d}_{:04d}.png'.format( hparams.sample_dir, epoch, batch_num)) if epoch % hparams.summary_epoch == 0: print "Epoch:", '%04d' % (epoch), 'Avg loss = {:.9f}'.format( avg_loss) if epoch % hparams.ckpt_epoch == 0: save_path = os.path.join(hparams.ckpt_dir, 'omniglot_vae_model') model_saver.save(sess, save_path, global_step=epoch) save_path = os.path.join(hparams.ckpt_dir, 'omniglot_vae_model') model_saver.save(sess, save_path, global_step=hparams.training_epochs - 1)
def stage_i(A_val,y_batch_val,hparams,hid_i,init_obj,early_stop,bs,optim,recovered=False): model_def = globals()['model_def'] m_loss1_batch_dict = {} m_loss2_batch_dict = {} zp_loss_batch_dict = {} total_loss_dict = {} x_hat_batch_dict = {} model_selection = ModelSelect(hparams) hid_i=int(hid_i) # print('Matrix norm is {}'.format(np.linalg.norm(A_val))) # hparams.eps = hparams.eps * np.linalg.norm(A_val) # Get a session sess = tf.Session() # Set up palceholders A = tf.placeholder(tf.float32, shape=(hparams.n_input, hparams.num_measurements), name='A') y_batch = tf.placeholder(tf.float32, shape=(hparams.batch_size, hparams.num_measurements), name='y_batch') # Create the generator model_hparams = model_def.Hparams() model_hparams.n_z = hparams.n_z model_hparams.stdv = hparams.stdv model_hparams.mean = hparams.mean model_hparams.grid = copy.deepcopy(hparams.grid) model_selection.setup_dim(hid_i,model_hparams) if not hparams.model_types[0] == 'vae-flex-alt' and 'alt' in hparams.model_types[0]: model_def.ignore_grid = next((j for j in model_selection.dim_list if j >= hid_i), None) #set up the initialization print('The initialization is: {}'.format(init_obj.mode)) if init_obj.mode=='random': z_batch = model_def.get_z_var(model_hparams,hparams.batch_size,hid_i) elif init_obj.mode in ['previous-and-random','only-previous']: z_batch = model_def.get_z_var(model_hparams,hparams.batch_size,hid_i) init_op_par = tf.assign(z_batch, truncate_val(model_hparams,hparams,hid_i,init_obj,stdv=0)) else: z_batch = truncate_val(model_hparams,hparams,hid_i,init_obj,stdv=0.1) _, x_hat_batch, _ = model_def.generator_i(model_hparams, z_batch, 'gen', hparams.bol,hid_i,relative=False) x_hat_batch_dict[hid_i] = x_hat_batch # measure the estimate if hparams.measurement_type == 'project': y_hat_batch = tf.identity(x_hat_batch, name='y_hat_batch') else: y_hat_batch = tf.matmul(x_hat_batch, A, name='y_hat_batch') # define all losses m_loss1_batch = tf.reduce_mean(tf.abs(y_batch - y_hat_batch), 1) m_loss2_batch = tf.reduce_mean((y_batch - y_hat_batch)**2, 1) if hparams.stdv>0: norm_val = 1/(hparams.stdv**2) else: norm_val = 1e+20 zp_loss_batch = tf.reduce_sum((z_batch-tf.ones(tf.shape(z_batch))*hparams.mean)**2*norm_val, 1) #added normalization # define total loss total_loss_batch = hparams.mloss1_weight * m_loss1_batch \ + hparams.mloss2_weight * m_loss2_batch \ + hparams.zprior_weight * zp_loss_batch total_loss = tf.reduce_mean(total_loss_batch) total_loss_dict[hid_i] = total_loss # Compute means for logging m_loss1 = tf.reduce_mean(m_loss1_batch) m_loss2 = tf.reduce_mean(m_loss2_batch) zp_loss = tf.reduce_mean(zp_loss_batch) m_loss1_batch_dict[hid_i] = m_loss1 m_loss2_batch_dict[hid_i] = m_loss2 zp_loss_batch_dict[hid_i] = zp_loss # Set up gradient descent var_list = [z_batch] if recovered: global_step = tf.Variable(hparams.optim.global_step, trainable=False, name='global_step') else: global_step = tf.Variable(0, trainable=False, name='global_step') learning_rate = utils.get_learning_rate(global_step, hparams) opt = utils.get_optimizer(learning_rate, hparams) update_op = opt.minimize(total_loss, var_list=var_list, global_step=global_step, name='update_op') opt_reinit_op = utils.get_opt_reinit_op(opt, var_list, global_step) # Intialize and restore model parameters init_op = tf.global_variables_initializer() sess.run(init_op) #restore the setting if 'alt' in hparams.model_types[0]: factor = 1 else: factor = len(hparams.grid) model_def.batch_size = hparams.batch_size*factor #changes object (call by reference), necessary, since call of generator_i might change batch size. model_selection.restore(sess,hid_i) if recovered: best_keeper = hparams.optim.best_keeper else: best_keeper = utils.BestKeeper(hparams,logg_z=True) if hparams.measurement_type == 'project': feed_dict = {y_batch: y_batch_val} else: feed_dict = {A: A_val, y_batch: y_batch_val} flag = False for i in range(init_obj.num_random_restarts): if recovered and i <= hparams.optim.i: #Loosing optimizer's state, keras implementation maybe better if i < hparams.optim.i: continue else: sess.run(utils.get_opt_reinit_op(opt, [], global_step)) sess.run(tf.assign(z_batch,hparams.optim.z_batch)) else: sess.run(opt_reinit_op) if i<1 and init_obj.mode in ['previous-and-random','only-previous']: print('Using previous outcome as starting point') sess.run(init_op_par) for j in range(hparams.max_update_iter): if recovered and j < hparams.optim.j: continue _, lr_val, total_loss_val, \ m_loss1_val, \ m_loss2_val, \ zp_loss_val = sess.run([update_op, learning_rate, total_loss, m_loss1, m_loss2, zp_loss], feed_dict=feed_dict) if hparams.gif and ((j % hparams.gif_iter) == 0): images = sess.run(x_hat_batch, feed_dict=feed_dict) for im_num, image in enumerate(images): save_dir = '{0}/{1}/{2}/'.format(hparams.gif_dir, hid_i,im_num) utils.set_up_dir(save_dir) save_path = save_dir + '{0}.png'.format(j) image = image.reshape(hparams.image_shape) save_image(image, save_path) if j%100==0 and early_stop: x_hat_batch_val = sess.run(x_hat_batch, feed_dict=feed_dict) if check_tolerance(hparams,A_val,x_hat_batch_val,y_batch_val)[1]: flag = True print('Early stopping') break if j%25==0:#Now not every turn logging_format = 'hid {} rr {} iter {} lr {} total_loss {} m_loss1 {} m_loss2 {} zp_loss {}' print( logging_format.format(hid_i, i, j, lr_val, total_loss_val, m_loss1_val, m_loss2_val, zp_loss_val)) if j%100==0: x_hat_batch_val, total_loss_batch_val, z_batch_val = sess.run([x_hat_batch, total_loss_batch,z_batch], feed_dict=feed_dict) best_keeper.report(x_hat_batch_val, total_loss_batch_val,z_val=z_batch_val) optim.global_step = sess.run(global_step) optim.A = A_val optim.y_batch = y_batch_val optim.i=i optim.j=j optim.z_batch= z_batch_val optim.best_keeper=best_keeper optim.bs=bs optim.init_obj = init_obj utils.save_to_pickle(optim,utils.get_checkpoint_dir(hparams, hparams.model_types[0])+'tmp/optim.pkl') print('Checkpoint of optimization created') hparams.optim.j = 0 x_hat_batch_val, total_loss_batch_val, z_batch_val = sess.run([x_hat_batch, total_loss_batch,z_batch], feed_dict=feed_dict) best_keeper.report(x_hat_batch_val, total_loss_batch_val,z_val=z_batch_val) if flag: break tf.reset_default_graph() return best_keeper.get_best()
def main(hparams): # os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # Set up some stuff according to hparams utils.set_up_dir(hparams.ckpt_dir) utils.set_up_dir(hparams.sample_dir) utils.print_hparams(hparams) # encode x_ph = tf.placeholder(tf.float32, [None, hparams.n_input], name='x_ph') _,x_reconstr_mean,_, loss_list = model_def.model(hparams,x_ph,['enc','gen'],[False,False]) total_loss = tf.add_n(loss_list, name='total_loss') # z_mean, z_log_sigma_sq = model_def.encoder(hparams, x_ph, 'enc', reuse=False) # # # sample # eps = tf.random_normal((hparams.batch_size, hparams.n_z), 0, 1, dtype=tf.float32) # z_sigma = tf.sqrt(tf.exp(z_log_sigma_sq)) # z = z_mean + z_sigma * eps # # # reconstruct # logits, x_reconstr_mean, _ = model_def.generator(hparams, z, 'gen', reuse=False) # # generator sampler z_ph = tf.placeholder(tf.float32, [None, hparams.grid[-1]], name='x_ph') x_sample = [] for i in range(len(hparams.grid)): _, x_sample_tmp, _ = model_def.generator_i(hparams, model_def.slicer_dec(hparams,i,None,z_ph), 'gen', True,i) x_sample.append(x_sample_tmp) # _, x_sample, _ = model_def.generator(hparams, z_ph, 'gen', reuse=True) # # define loss and update op # total_loss = model_def.get_loss(x_ph, logits, z_mean, z_log_sigma_sq) opt = tf.train.AdamOptimizer(learning_rate=hparams.learning_rate) update_op = opt.minimize(total_loss) # print([el.name for el in tf.trainable_variables()]) # Sanity checks for var in tf.global_variables(): print(var.op.name) print('') # print([o.name for o in tf.trainable_variables()]) # Get a new session sess = tf.Session() # Model checkpointing setup model_saver = tf.train.Saver() init_op = tf.global_variables_initializer() sess.run(init_op) # Attempt to restore variables from checkpoint start_epoch = utils.try_restore(hparams, sess, model_saver) # Get data iterator iterator = data_input.mnist_data_iteratior(dataset=hparams.dataset) # Training for epoch in range(start_epoch+1, hparams.training_epochs): avg_loss = 0.0 num_batches = hparams.num_samples // hparams.batch_size batch_num = 0 for (x_batch_val, _) in iterator(hparams, num_batches): batch_num += 1 feed_dict = {x_ph: x_batch_val} _, loss_val = sess.run([update_op, total_loss], feed_dict=feed_dict) avg_loss += loss_val / hparams.num_samples * hparams.batch_size if batch_num % 100 == 0: x_reconstr_mean_val = sess.run(x_reconstr_mean, feed_dict={x_ph: x_batch_val}) z_val = np.random.randn(hparams.batch_size, hparams.grid[-1]) x_sample_val = sess.run(x_sample, feed_dict={z_ph: z_val}) # print(sess.run(hparams.track[0], feed_dict={z_ph: z_val})) # print(sess.run(hparams.track[1], feed_dict={z_ph: z_val})) # s1 = sess.run(hparams.x_ph_ref[0], feed_dict={x_ph: x_batch_val}) # s2 = sess.run(hparams.logits_ref[0], feed_dict={x_ph: x_batch_val}) # s3 = sess.run(loss_list, feed_dict={x_ph: x_batch_val}) # print(s1.shape) # print(s2.shape) # print(s3) utils.save_images(np.reshape(x_batch_val, [-1, 28, 28]), \ [10, 10], \ '{}/orig_{:02d}_{:04d}.png'.format(hparams.sample_dir, epoch, batch_num)) for i in range(len(hparams.grid)): utils.save_images(np.reshape(x_reconstr_mean_val[i], [-1, 28, 28]), [10, 10], '{}/reconstr_{}_{:02d}_{:04d}.png'.format(hparams.sample_dir, hparams.grid[i], epoch, batch_num)) utils.save_images(np.reshape(x_sample_val[i], [-1, 28, 28]), [10, 10], '{}/sampled_{}_{:02d}_{:04d}.png'.format(hparams.sample_dir, hparams.grid[i], epoch, batch_num)) if epoch % hparams.summary_epoch == 0: print("Epoch:", '%04d' % (epoch), 'Avg loss = {:.9f}'.format(avg_loss)) if epoch % hparams.ckpt_epoch == 0: save_path = os.path.join(hparams.ckpt_dir, '{}_vae_model_flex_hid'.format(hparams.dataset)+str('_'.join(map(str,hparams.grid)))) model_saver.save(sess, save_path, global_step=epoch) save_path = os.path.join(hparams.ckpt_dir, '{}_vae_model_flex_hid'.format(hparams.dataset)+str('_'.join(map(str,hparams.grid)))) model_saver.save(sess, save_path, global_step=hparams.training_epochs-1)
} }, 'N50NUR': {'template_net': "/mnt/antares_raid/home/oliver/Scripts/template_AE_net_in_out.prototxt" , 'template_solver': "/mnt/antares_raid/home/oliver/Scripts/template_autoencoder_solver.prototxt", 'replacement_dict': { 'train_net': '"{}"'.format(lmdb_root + "lmdb_corrupted_50/MNIST_TRAIN_60000_corrupt_px_392_unrot_corrupted_lmdb/shuffled/"), 'test_net': '"{}"'.format(lmdb_root + "lmdb_corrupted_50/MNIST_TEST_10000_corrupt_px_392_unrot_corrupted_lmdb/shuffled/" ), 'train_net_out': '"{}"'.format(lmdb_root + "lmdb_corrupted_50/MNIST_TRAIN_60000_corrupt_px_392_unrot_lmdb/shuffled/"), 'test_net_out': '"{}"'.format(lmdb_root + "lmdb_corrupted_50/MNIST_TEST_10000_corrupt_px_392_unrot_lmdb/shuffled/") } } } set_up_dir(init_weights) for c in cases: replacement_dict_solver = { 'network': [], 'testiter': 130, 'testinterval': 780, 'maxiter': 390000, 'snapsht': 10000, 'snapshotprefix': [] } print(c) set_up_dir(root + c + '/snapshots/') set_up_dir(root + c + '/log/') net_file = root + c + '/net.prototxt'