def ckpt_pb(pb_path): tf.reset_default_graph() config = tf.ConfigProto(allow_soft_placement=True) config.gpu_options.allow_growth = True with tf.Session() as sess: path_base = Path(g_config['eval_dataset']).parent.parent _mean_shape = mio.import_pickle(path_base / 'reference_shape.pkl') _mean_shape = data_provider.align_reference_shape_to_112(_mean_shape) assert isinstance(_mean_shape, np.ndarray) print(_mean_shape.shape) tf_img = tf.placeholder(dtype=tf.float32, shape=(1, 112, 112, 3), name='inputs/input_img') tf_dummy = tf.placeholder(dtype=tf.float32, shape=(1, 73, 2), name='inputs/input_shape') tf_shape = tf.constant(_mean_shape, dtype=tf.float32, shape=(73, 2), name='MeanShape') model = mdm_model.MDMModel(tf_img, tf_dummy, tf_shape, batch_size=1, num_patches=g_config['num_patches'], num_channels=3, is_training=False) saver = tf.train.Saver() ckpt = tf.train.get_checkpoint_state(g_config['train_dir']) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) global_step = ckpt.model_checkpoint_path.split('/')[-1].split( '-')[-1] print('Successfully loaded model from {} at step={}.'.format( ckpt.model_checkpoint_path, global_step)) else: print('No checkpoint file found') return output_graph_def = tf.graph_util.convert_variables_to_constants( sess, tf.get_default_graph().as_graph_def(), ['Network/Predict/add']) with tf.gfile.FastGFile(pb_path, mode='wb') as f: f.write(output_graph_def.SerializeToString())
def train(scope=''): """Train on dataset for a number of steps.""" with tf.Graph().as_default() as graph, tf.device('/gpu:0'): # Global steps tf_global_step = tf.get_variable( 'GlobalStep', [], initializer=tf.constant_initializer(0), trainable=False) # Learning rate tf_lr = tf.train.exponential_decay(g_config['learning_rate'], tf_global_step, g_config['learning_rate_step'], g_config['learning_rate_decay'], staircase=True, name='LearningRate') tf.summary.scalar('learning_rate', tf_lr, collections=['train']) # Create an optimizer that performs gradient descent. opt = tf.train.AdamOptimizer(tf_lr) data_provider.prepare_images(g_config['train_dataset'].split(':'), num_patches=g_config['num_patches'], verbose=True) path_base = Path(g_config['train_dataset'].split(':')[0]).parent.parent _mean_shape = mio.import_pickle(path_base / 'reference_shape.pkl') _mean_shape = data_provider.align_reference_shape_to_112(_mean_shape) assert (isinstance(_mean_shape, np.ndarray)) assert (_mean_shape.shape[0] == g_config['num_patches']) tf_mean_shape = tf.constant(_mean_shape, dtype=tf.float32, name='MeanShape') def get_random_sample(image, shape, rotation_stddev=10): # Read a random image with landmarks and bb image = menpo.image.Image(image.transpose((2, 0, 1)), copy=False) image.landmarks['PTS'] = PointCloud(shape) if np.random.rand() < .5: image = utils.mirror_image(image) if np.random.rand() < .5: theta = np.random.normal(scale=rotation_stddev) rot = menpo.transform.rotate_ccw_about_centre( image.landmarks['PTS'], theta) image = image.warp_to_shape(image.shape, rot) bb = image.landmarks['PTS'].bounding_box().points miny, minx = np.min(bb, 0) maxy, maxx = np.max(bb, 0) bbsize = max(maxx - minx, maxy - miny) center = [(miny + maxy) / 2., (minx + maxx) / 2.] shift = (np.random.rand(2) - 0.5) * 0.6 * bbsize image.landmarks['bb'] = PointCloud([ [ center[0] - bbsize * 0.5 + shift[0], center[1] - bbsize * 0.5 + shift[1] ], [ center[0] + bbsize * 0.5 + shift[0], center[1] + bbsize * 0.5 + shift[1] ], ]).bounding_box() proportion = 1.0 / 6.0 + float(np.random.rand() - 0.5) / 6. image = image.crop_to_landmarks_proportion(proportion, group='bb') image = image.resize((112, 112)) random_image = image.pixels.transpose(1, 2, 0).astype('float32') random_shape = image.landmarks['PTS'].points.astype('float32') return random_image, random_shape def decode_feature_and_augment(serialized): feature = { 'train/image': tf.FixedLenFeature([], tf.string), 'train/shape': tf.VarLenFeature(tf.float32), } features = tf.parse_single_example(serialized, features=feature) decoded_image = tf.decode_raw(features['train/image'], tf.float32) decoded_image = tf.reshape(decoded_image, (448, 448, 3)) decoded_shape = tf.sparse.to_dense(features['train/shape']) decoded_shape = tf.reshape(decoded_shape, (g_config['num_patches'], 2)) random_image, random_shape = tf.py_func( get_random_sample, [decoded_image, decoded_shape], [tf.float32, tf.float32], stateful=True, name='RandomSample') return data_provider.distort_color(random_image), random_shape def decode_feature(serialized): feature = { 'validate/image': tf.FixedLenFeature([], tf.string), 'validate/shape': tf.VarLenFeature(tf.float32), } features = tf.parse_single_example(serialized, features=feature) decoded_image = tf.decode_raw(features['validate/image'], tf.float32) decoded_image = tf.reshape(decoded_image, (112, 112, 3)) decoded_shape = tf.sparse.to_dense(features['validate/shape']) decoded_shape = tf.reshape(decoded_shape, (g_config['num_patches'], 2)) return decoded_image, decoded_shape with tf.name_scope('DataProvider'): tf_dataset = tf.data.TFRecordDataset( [str(path_base / 'train.bin')]) tf_dataset = tf_dataset.repeat() tf_dataset = tf_dataset.map(decode_feature_and_augment, num_parallel_calls=5) tf_dataset = tf_dataset.batch(g_config['batch_size'], True) tf_dataset = tf_dataset.prefetch(1) tf_iterator = tf_dataset.make_one_shot_iterator() tf_images, tf_shapes = tf_iterator.get_next(name='Batch') tf_images.set_shape([g_config['batch_size'], 112, 112, 3]) tf_shapes.set_shape([g_config['batch_size'], 73, 2]) tf_dataset_v = tf.data.TFRecordDataset( [str(path_base / 'validate.bin')]) tf_dataset_v = tf_dataset_v.repeat() tf_dataset_v = tf_dataset_v.map(decode_feature, num_parallel_calls=5) tf_dataset_v = tf_dataset_v.batch(50, True) tf_dataset_v = tf_dataset_v.prefetch(1) tf_iterator_v = tf_dataset_v.make_one_shot_iterator() tf_images_v, tf_shapes_v = tf_iterator_v.get_next( name='ValidateBatch') tf_images_v.set_shape([50, 112, 112, 3]) tf_shapes_v.set_shape([50, 73, 2]) print('Defining model...') with tf.device(g_config['train_device']): tf_model = mdm_model.MDMModel(tf_images, tf_shapes, tf_mean_shape, batch_size=g_config['batch_size'], num_patches=g_config['num_patches'], num_channels=3) tf_grads = opt.compute_gradients(tf_model.nme) with tf.name_scope('Validate'): tf_model_v = mdm_model.MDMModel( tf_images_v, tf_shapes_v, tf_mean_shape, batch_size=50, num_patches=g_config['num_patches'], num_channels=3, is_training=False) tf.summary.histogram('dx', tf_model.prediction - tf_shapes, collections=['train']) bn_updates = tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope) # Add histograms for gradients. for grad, var in tf_grads: if grad is not None: tf.summary.histogram(var.op.name + '/gradients', grad, collections=['train']) # Apply the gradients to adjust the shared variables. with tf.name_scope('Optimizer', values=[tf_grads, tf_global_step]): apply_gradient_op = opt.apply_gradients(tf_grads, global_step=tf_global_step) # Add histograms for trainable variables. for var in tf.trainable_variables(): tf.summary.histogram(var.op.name, var, collections=['train']) with tf.name_scope('MovingAverage', values=[tf_global_step]): variable_averages = tf.train.ExponentialMovingAverage( g_config['MOVING_AVERAGE_DECAY'], tf_global_step) variables_to_average = (tf.trainable_variables() + tf.moving_average_variables()) variables_averages_op = variable_averages.apply( variables_to_average) # Group all updates to into a single train op. bn_updates_op = tf.group(*bn_updates, name='BNGroup') train_op = tf.group(apply_gradient_op, variables_averages_op, bn_updates_op, name='TrainGroup') # Create a saver. saver = tf.train.Saver() train_summary_op = tf.summary.merge_all('train') validate_summary_op = tf.summary.merge_all('validate') config = tf.ConfigProto(allow_soft_placement=True) config.gpu_options.allow_growth = True sess = tf.Session(graph=graph, config=config) init = tf.global_variables_initializer() print('Initializing variables...') sess.run(init) print('Initialized variables.') # Assuming model_checkpoint_path looks something like: # /ckpt/train/model.ckpt-0, # extract global_step from it. start_step = 0 ckpt = tf.train.get_checkpoint_state(g_config['train_dir']) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) start_step = int( ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]) + 1 print('%s: Restart from %s' % (datetime.now(), g_config['train_dir'])) else: ckpt = tf.train.get_checkpoint_state(g_config['ckpt_dir']) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) tf_global_step_op = tf_global_step.assign(0) sess.run(tf_global_step_op) print('%s: Pre-trained model restored from %s' % (datetime.now(), g_config['ckpt_dir'])) train_writer = tf.summary.FileWriter(g_config['train_dir'] + '/train', sess.graph) validate_writer = tf.summary.FileWriter( g_config['train_dir'] + '/validate', sess.graph) print('Starting training...') steps_per_epoch = 15000 / g_config['batch_size'] for step in range(start_step, g_config['max_steps']): if step % steps_per_epoch == 0: start_time = time.time() _, train_loss, train_summary = sess.run( [train_op, tf_model.nme, train_summary_op]) duration = time.time() - start_time validate_loss, validate_summary = sess.run( [tf_model_v.nme, validate_summary_op]) train_writer.add_summary(train_summary, step) validate_writer.add_summary(validate_summary, step) print('%s: step %d, loss = %.4f (%.3f sec/batch)' % (datetime.now(), step, train_loss, duration)) print('%s: step %d, validate loss = %.4f' % (datetime.now(), step, validate_loss)) else: start_time = time.time() _, train_loss = sess.run([train_op, tf_model.nme]) duration = time.time() - start_time if step % 100 == 0: print('%s: step %d, loss = %.4f (%.3f sec/batch)' % (datetime.now(), step, train_loss, duration)) assert not np.isnan(train_loss), 'Model diverged with loss = NaN' if step % steps_per_epoch == 0 or (step + 1) == g_config['max_steps']: checkpoint_path = os.path.join(g_config['train_dir'], 'model.ckpt') saver.save(sess, checkpoint_path, global_step=step)
def evaluate(): with tf.Graph().as_default(), tf.device('/cpu:0'): path_base = Path(g_config['eval_dataset']).parent.parent _mean_shape = mio.import_pickle(path_base / 'reference_shape.pkl') _mean_shape = data_provider.align_reference_shape_to_112(_mean_shape) tf_mean_shape = tf.constant(_mean_shape, dtype=tf.float32, name='MeanShape') def decode_feature(serialized): feature = { 'test/image': tf.FixedLenFeature([], tf.string), 'test/shape': tf.VarLenFeature(tf.float32), 'test/init': tf.VarLenFeature(tf.float32), } features = tf.parse_single_example(serialized, features=feature) decoded_image = tf.decode_raw(features['test/image'], tf.float32) decoded_image = tf.reshape(decoded_image, (112, 112, 3)) decoded_shape = tf.sparse.to_dense(features['test/shape']) decoded_shape = tf.reshape(decoded_shape, (g_config['num_patches'], 2)) return decoded_image, decoded_shape with tf.name_scope('DataProvider', values=[]): tf_dataset = tf.data.TFRecordDataset([str(path_base / 'test.bin')]) tf_dataset = tf_dataset.map(decode_feature) tf_dataset = tf_dataset.batch(1) tf_dataset = tf_dataset.prefetch(1000) tf_iterator = tf_dataset.make_one_shot_iterator() tf_images, tf_shapes = tf_iterator.get_next(name='batch') tf_images.set_shape((1, 112, 112, 3)) tf_shapes.set_shape((1, 73, 2)) print('Loading model...') with tf.device(g_config['eval_device']): model = mdm_model.MDMModel(tf_images, tf_shapes, tf_mean_shape, batch_size=1, num_patches=g_config['num_patches'], num_channels=3, is_training=False) # Restore the moving average version of the learned variables for eval. variable_averages = tf.train.ExponentialMovingAverage( g_config['MOVING_AVERAGE_DECAY']) variables_to_restore = variable_averages.variables_to_restore() saver = tf.train.Saver(variables_to_restore) graph_def = tf.get_default_graph().as_graph_def() summary_writer = tf.summary.FileWriter(g_config['eval_dir'], graph_def=graph_def) config = tf.ConfigProto(allow_soft_placement=True) config.gpu_options.allow_growth = True with tf.Session(config=config) as sess: ckpt = tf.train.get_checkpoint_state(g_config['train_dir']) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) # Assuming model_checkpoint_path looks something like: # /ckpt/train/model.ckpt-0, # extract global_step from it. global_step = ckpt.model_checkpoint_path.split('/')[-1].split( '-')[-1] print('Successfully loaded model from {} at step={}.'.format( ckpt.model_checkpoint_path, global_step)) else: print('No checkpoint file found') return eval_base = Path('Evaluate') for i in range(10): eval_path = eval_base / 'err{}'.format(i) if not eval_path.exists(): eval_path.mkdir(parents=True) num_iter = g_config['num_examples'] # Counts the number of correct predictions. errors = [] mean_errors = [] print('%s: starting evaluation on (%s).' % (datetime.now(), g_config['eval_dataset'])) start_time = time.time() for step in range(num_iter): nme, ne, img = sess.run( [model.batch_nme, model.batch_ne, model.out_images]) error_level = min(9, int(nme[0] * 100)) plt.imsave( 'Evaluate/err{}/step{}.png'.format(error_level, step), img[0]) errors.append(ne) mean_errors.append(nme) step += 1 if step % 20 == 0: duration = time.time() - start_time sec_per_batch = duration / 20.0 examples_per_sec = 1. / sec_per_batch log_str = '{}: [{:d} batches out of {:d}] ({:.1f} examples/sec; {:.3f} sec/batch)' print( log_str.format(datetime.now(), step, num_iter, examples_per_sec, sec_per_batch)) start_time = time.time() errors = np.array(errors) errors = np.reshape(errors, (-1, g_config['num_patches'])) print(errors.shape) mean_errors = np.vstack(mean_errors).ravel() mean_rse = np.mean(errors, 0) mean_rmse = mean_errors.mean() with open('Evaluate/errors.txt', 'w') as ofs: for row, avg in zip(errors, mean_errors): for col in row: ofs.write('%.4f, ' % col) ofs.write('%.4f' % avg) ofs.write('\n') for col in mean_rse: ofs.write('%.4f, ' % col) ofs.write('%.4f' % mean_rmse) ofs.write('\n') auc_at_08 = (mean_errors < .08).mean() auc_at_05 = (mean_errors < .05).mean() print('Errors', mean_errors.shape) print( '%s: mean_rmse = %.4f, auc @ 0.05 = %.4f, auc @ 0.08 = %.4f [%d examples]' % (datetime.now(), mean_errors.mean(), auc_at_05, auc_at_08, num_iter)) ced_image = plot_ced([mean_errors.tolist()], ['MDM']) ced_plot = sess.run( tf.summary.merge( [tf.summary.image('ced_plot', ced_image[None, ...])])) summary_writer.add_summary(ced_plot, global_step)
def train(scope=''): """Train on dataset for a number of steps.""" with tf.Graph().as_default() as graph, tf.device('/gpu:0'): # Global steps tf_global_step = tf.get_variable( 'GlobalStep', [], initializer=tf.constant_initializer(0), trainable=False) # Learning rate tf_lr = tf.train.exponential_decay(g_config['learning_rate'], tf_global_step, g_config['learning_rate_step'], g_config['learning_rate_decay'], staircase=True, name='LearningRate') tf.summary.scalar('learning_rate', tf_lr) # Create an optimizer that performs gradient descent. opt = tf.train.AdamOptimizer(tf_lr) data_provider.prepare_images(g_config['train_dataset'].split(':'), num_patches=g_config['num_patches'], verbose=True) path_base = Path(g_config['train_dataset'].split(':')[0]).parent.parent _mean_shape = mio.import_pickle(path_base / 'reference_shape.pkl') _mean_shape = data_provider.align_reference_shape_to_112(_mean_shape) assert (isinstance(_mean_shape, np.ndarray)) assert (_mean_shape.shape[0] == g_config['num_patches']) tf_mean_shape = tf.constant(_mean_shape, dtype=tf.float32, name='MeanShape') def decode_feature(serialized): feature = { 'train/image': tf.FixedLenFeature([], tf.string), 'train/shape': tf.VarLenFeature(tf.float32), } features = tf.parse_single_example(serialized, features=feature) decoded_image = tf.decode_raw(features['train/image'], tf.float32) decoded_image = tf.reshape(decoded_image, (336, 336, 3)) decoded_shape = tf.sparse.to_dense(features['train/shape']) decoded_shape = tf.reshape(decoded_shape, (g_config['num_patches'], 2)) return decoded_image, decoded_shape def get_random_sample(image, shape, rotation_stddev=10): # Read a random image with landmarks and bb image = menpo.image.Image(image.transpose((2, 0, 1)), copy=False) image.landmarks['PTS'] = PointCloud(shape) if np.random.rand() < .5: image = utils.mirror_image(image) if np.random.rand() < .5: theta = np.random.normal(scale=rotation_stddev) rot = menpo.transform.rotate_ccw_about_centre( image.landmarks['PTS'], theta) image = image.warp_to_shape(image.shape, rot) bb = image.landmarks['PTS'].bounding_box().points miny, minx = np.min(bb, 0) maxy, maxx = np.max(bb, 0) bbsize = max(maxx - minx, maxy - miny) center = [(miny + maxy) / 2., (minx + maxx) / 2.] shift = (np.random.rand(2) - 0.5) / 3. * bbsize image.landmarks['bb'] = PointCloud([ [ center[0] - bbsize * 0.5 + shift[0], center[1] - bbsize * 0.5 + shift[1] ], [ center[0] + bbsize * 0.5 + shift[0], center[1] + bbsize * 0.5 + shift[1] ], ]).bounding_box() proportion = 1.0 / 6.0 + float(np.random.rand() - 0.5) / 50.0 image = image.crop_to_landmarks_proportion(proportion, group='bb') image = image.resize((112, 112)) random_image = image.pixels.transpose(1, 2, 0).astype('float32') random_shape = image.landmarks['PTS'].points.astype('float32') return random_image, random_shape def get_init_shape(image, shape): return image, shape, tf_mean_shape def distort_color(image, shape, init_shape): return data_provider.distort_color(image), shape, init_shape with tf.name_scope('DataProvider', values=[tf_mean_shape]): tf_dataset = tf.data.TFRecordDataset( [str(path_base / 'train.bin')]) tf_dataset = tf_dataset.repeat() tf_dataset = tf_dataset.map(decode_feature) tf_dataset = tf_dataset.map(lambda x, y: tf.py_func( get_random_sample, [x, y], [tf.float32, tf.float32], stateful=True, name='RandomSample')) tf_dataset = tf_dataset.map(get_init_shape) tf_dataset = tf_dataset.map(distort_color) tf_dataset = tf_dataset.batch(g_config['batch_size'], True) tf_dataset = tf_dataset.prefetch(7500) tf_iterator = tf_dataset.make_one_shot_iterator() tf_images, tf_shapes, tf_initial_shapes = tf_iterator.get_next( name='Batch') tf_images.set_shape([g_config['batch_size'], 112, 112, 3]) tf_shapes.set_shape([g_config['batch_size'], 73, 2]) tf_initial_shapes.set_shape([g_config['batch_size'], 73, 2]) print('Defining model...') with tf.device(g_config['train_device']): tf_model = mdm_model.MDMModel( tf_images, tf_shapes, tf_initial_shapes, batch_size=g_config['batch_size'], num_iterations=g_config['num_iterations'], num_patches=g_config['num_patches'], patch_shape=(g_config['patch_size'], g_config['patch_size']), num_channels=3) with tf.name_scope('Losses', values=[tf_model.prediction, tf_shapes]): tf_norm_error = tf_model.normalized_rmse( tf_model.prediction, tf_shapes) tf_loss = tf.reduce_mean(tf_norm_error) tf.summary.scalar('losses/total', tf_loss) # Calculate the gradients for the batch of data tf_grads = opt.compute_gradients(tf_loss) tf.summary.histogram('dx', tf_model.prediction - tf_shapes) bn_updates = tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope) # Add histograms for gradients. for grad, var in tf_grads: if grad is not None: tf.summary.histogram(var.op.name + '/gradients', grad) # Apply the gradients to adjust the shared variables. with tf.name_scope('Optimizer', values=[tf_grads, tf_global_step]): apply_gradient_op = opt.apply_gradients(tf_grads, global_step=tf_global_step) # Add histograms for trainable variables. for var in tf.trainable_variables(): tf.summary.histogram(var.op.name, var) with tf.name_scope('MovingAverage', values=[tf_global_step]): variable_averages = tf.train.ExponentialMovingAverage( g_config['MOVING_AVERAGE_DECAY'], tf_global_step) variables_to_average = (tf.trainable_variables() + tf.moving_average_variables()) variables_averages_op = variable_averages.apply( variables_to_average) # Group all updates to into a single train op. bn_updates_op = tf.group(*bn_updates, name='BNGroup') train_op = tf.group(apply_gradient_op, variables_averages_op, bn_updates_op, name='TrainGroup') # Create a saver. saver = tf.train.Saver() # Build the summary operation from the last tower summaries. summary_op = tf.summary.merge_all() # Start running operations on the Graph. allow_soft_placement must be # set to True to build towers on GPU, as some of the ops do not have GPU # implementations. config = tf.ConfigProto(allow_soft_placement=True) config.gpu_options.allow_growth = True sess = tf.Session(config=config) # Build an initialization operation to run below. init = tf.global_variables_initializer() print('Initializing variables...') sess.run(init) print('Initialized variables.') start_step = 0 ckpt = tf.train.get_checkpoint_state(g_config['train_dir']) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) # Assuming model_checkpoint_path looks something like: # /ckpt/train/model.ckpt-0, # extract global_step from it. start_step = int( ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]) + 1 print('%s: Pre-trained model restored from %s' % (datetime.now(), g_config['train_dir'])) elif TUNE: assign_op = [] vvv = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) cnt = 0 for v in vvv: if v.name in tf_model.var_map: assign_op.append( v.assign( graph.get_tensor_by_name( tf_model.var_map[v.name]))) cnt += 1 else: print(v.name) sess.run(assign_op) print('%s: Pre-trained model restored from graph.pb %d/%d' % (datetime.now(), cnt, len(vvv))) summary_writer = tf.summary.FileWriter(g_config['train_dir'], sess.graph) print('Starting training...') for step in range(start_step, g_config['max_steps']): start_time = time.time() _, loss_value = sess.run([train_op, tf_loss]) duration = time.time() - start_time assert not np.isnan(loss_value), 'Model diverged with loss = NaN' if step % 100 == 0: examples_per_sec = g_config['batch_size'] / float(duration) format_str = ( '%s: step %d, loss = %.4f (%.1f examples/sec; %.3f ' 'sec/batch)') print(format_str % (datetime.now(), step, loss_value, examples_per_sec, duration)) if step % 200 == 0: summary_str = sess.run(summary_op) summary_writer.add_summary(summary_str, step) # Save the model checkpoint periodically. if step % 1000 == 0 or (step + 1) == g_config['max_steps']: checkpoint_path = os.path.join(g_config['train_dir'], 'model.ckpt') saver.save(sess, checkpoint_path, global_step=step)
def train(scope=''): """Train on dataset for a number of steps.""" with tf.Graph().as_default() as graph, tf.device('/gpu:0'): # Global steps tf_global_step = tf.get_variable( 'GlobalStep', [], initializer=tf.constant_initializer(0), trainable=False ) # Learning rate tf_lr = tf.train.exponential_decay( g_config['learning_rate'], tf_global_step, g_config['learning_rate_step'], g_config['learning_rate_decay'], staircase=True, name='LearningRate' ) tf.summary.scalar('learning_rate', tf_lr, collections=['train']) # Create an optimizer that performs gradient descent. opt = tf.train.AdamOptimizer(tf_lr) data_provider.prepare_images( g_config['train_dataset'].split(':'), num_patches=g_config['num_patches'], verbose=True ) path_base = Path(g_config['train_dataset'].split(':')[0]).parent.parent _mean_shape = mio.import_pickle(path_base / 'mean_shape.pkl') _mean_shape = data_provider.align_reference_shape_to_112(_mean_shape) assert(isinstance(_mean_shape, np.ndarray)) assert(_mean_shape.shape[0] == g_config['num_patches']) _negatives = [] for mp_image in mio.import_images('Dataset/Neg/*.png', verbose=True): _negatives.append(mp_image.pixels.transpose(1, 2, 0).astype(np.float32)) _num_negatives = len(_negatives) tf_mean_shape = tf.constant(_mean_shape, dtype=tf.float32, name='MeanShape') def get_random_sample(image, shape): # Occlude _O_AREA = 0.15 _O_MIN_H = 0.15 _O_MAX_H = 1.0 if np.random.rand() < .3: rh = min(112, int((np.random.rand() * (_O_MAX_H - _O_MIN_H) + _O_MIN_H) * 112)) rw = min(112, int(12544 * _O_AREA / rh)) dy = int(np.random.rand() * (112 - rh)) dx = int(np.random.rand() * (112 - rw)) idx = int(np.random.rand() * _num_negatives) image[dy:dy+rh, dx:dx+rw] = np.minimum( 1.0, _negatives[idx][dy:dy+rh, dx:dx+rw] ) return image, shape def decode_feature_and_augment(serialized): feature = { 'train/image': tf.FixedLenFeature([], tf.string), 'train/shape': tf.VarLenFeature(tf.float32), } features = tf.parse_single_example(serialized, features=feature) decoded_image = tf.decode_raw(features['train/image'], tf.float32) decoded_image = tf.reshape(decoded_image, (112, 112, 3)) decoded_shape = tf.sparse.to_dense(features['train/shape']) decoded_shape = tf.reshape(decoded_shape, (g_config['num_patches'], 2)) #decoded_image, decoded_shape = tf.py_func( # get_random_sample, [decoded_image, decoded_shape], [tf.float32, tf.float32], # stateful=True, # name='RandomSample' #) return data_provider.distort_color(decoded_image), decoded_shape def decode_feature(serialized): feature = { 'validate/image': tf.FixedLenFeature([], tf.string), 'validate/shape': tf.VarLenFeature(tf.float32), } features = tf.parse_single_example(serialized, features=feature) decoded_image = tf.decode_raw(features['validate/image'], tf.float32) decoded_image = tf.reshape(decoded_image, (112, 112, 3)) decoded_shape = tf.sparse.to_dense(features['validate/shape']) decoded_shape = tf.reshape(decoded_shape, (g_config['num_patches'], 2)) return decoded_image, decoded_shape with tf.name_scope('DataProvider'): tf_dataset = tf.data.TFRecordDataset([ str(path_base / 'train_0.bin'), str(path_base / 'train_1.bin'), str(path_base / 'train_2.bin'), str(path_base / 'train_3.bin') ]) tf_dataset = tf_dataset.repeat() tf_dataset = tf_dataset.map(decode_feature_and_augment, num_parallel_calls=5) tf_dataset = tf_dataset.shuffle(480) tf_dataset = tf_dataset.batch(g_config['batch_size'], True) tf_dataset = tf_dataset.prefetch(1) tf_iterator = tf_dataset.make_one_shot_iterator() tf_images, tf_shapes = tf_iterator.get_next(name='Batch') tf_images.set_shape([g_config['batch_size'], 112, 112, 3]) tf_shapes.set_shape([g_config['batch_size'], 75, 2]) tf_dataset_v = tf.data.TFRecordDataset([str(path_base / 'validate.bin')]) tf_dataset_v = tf_dataset_v.repeat() tf_dataset_v = tf_dataset_v.map(decode_feature, num_parallel_calls=5) tf_dataset_v = tf_dataset_v.batch(50, True) tf_dataset_v = tf_dataset_v.prefetch(1) tf_iterator_v = tf_dataset_v.make_one_shot_iterator() tf_images_v, tf_shapes_v = tf_iterator_v.get_next(name='ValidateBatch') tf_images_v.set_shape([50, 112, 112, 3]) tf_shapes_v.set_shape([50, 75, 2]) print('Defining model...') with tf.device(g_config['train_device']): tf_model = mdm_model.MDMModel( tf_images, tf_shapes, tf_mean_shape, batch_size=g_config['batch_size'], num_patches=g_config['num_patches'], num_channels=3, multiplier=g_config['multiplier'] ) tf_grads = opt.compute_gradients(tf_model.nme) with tf.name_scope('Validate'): tf_model_v = mdm_model.MDMModel( tf_images_v, tf_shapes_v, tf_mean_shape, batch_size=50, num_patches=g_config['num_patches'], num_channels=3, multiplier=g_config['multiplier'], is_training=False ) tf.summary.histogram('dx', tf_model.prediction - tf_shapes, collections=['train']) bn_updates = tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope) # Add histograms for gradients. for grad, var in tf_grads: if grad is not None: tf.summary.histogram(var.op.name + '/gradients', grad, collections=['train']) # Apply the gradients to adjust the shared variables. with tf.name_scope('Optimizer', values=[tf_grads, tf_global_step]): apply_gradient_op = opt.apply_gradients(tf_grads, global_step=tf_global_step) # Add histograms for trainable variables. for var in tf.trainable_variables(): tf.summary.histogram(var.op.name, var, collections=['train']) with tf.name_scope('MovingAverage', values=[tf_global_step]): variable_averages = tf.train.ExponentialMovingAverage(g_config['MOVING_AVERAGE_DECAY'], tf_global_step) variables_to_average = (tf.trainable_variables() + tf.moving_average_variables()) variables_averages_op = variable_averages.apply(variables_to_average) # Group all updates to into a single train op. bn_updates_op = tf.group(*bn_updates, name='BNGroup') train_op = tf.group( apply_gradient_op, variables_averages_op, bn_updates_op, name='TrainGroup' ) # Create a saver. saver = tf.train.Saver() train_summary_op = tf.summary.merge_all('train') validate_summary_op = tf.summary.merge_all('validate') config = tf.ConfigProto(allow_soft_placement=True) config.gpu_options.allow_growth = True sess = tf.Session(graph=graph, config=config) init = tf.global_variables_initializer() print('Initializing variables...') sess.run(init) print('Initialized variables.') # Assuming model_checkpoint_path looks something like: # /ckpt/train/model.ckpt-0, # extract global_step from it. start_step = 0 ckpt = tf.train.get_checkpoint_state(g_config['train_dir']) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) start_step = int(ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]) + 1 print('%s: Restart from %s' % (datetime.now(), g_config['train_dir'])) else: ckpt = tf.train.get_checkpoint_state(g_config['ckpt_dir']) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) tf_global_step_op = tf_global_step.assign(0) sess.run(tf_global_step_op) print('%s: Pre-trained model restored from %s' % (datetime.now(), g_config['ckpt_dir'])) train_writer = tf.summary.FileWriter(g_config['train_dir'] + '/train', sess.graph) validate_writer = tf.summary.FileWriter(g_config['train_dir'] + '/validate', sess.graph) print('Starting training...') steps_per_epoch = 15000 / g_config['batch_size'] for step in range(start_step, g_config['max_steps']): if step % steps_per_epoch == 0: start_time = time.time() _, train_loss, train_summary = sess.run([train_op, tf_model.nme, train_summary_op]) duration = time.time() - start_time validate_loss, validate_summary = sess.run([tf_model_v.nme, validate_summary_op]) train_writer.add_summary(train_summary, step) validate_writer.add_summary(validate_summary, step) print( '%s: step %d, loss = %.4f (%.3f sec/batch)' % ( datetime.now(), step, train_loss, duration ) ) print( '%s: step %d, validate loss = %.4f' % ( datetime.now(), step, validate_loss ) ) else: start_time = time.time() _, train_loss = sess.run([train_op, tf_model.nme]) duration = time.time() - start_time if step % 100 == 0: print( '%s: step %d, loss = %.4f (%.3f sec/batch)' % ( datetime.now(), step, train_loss, duration ) ) assert not np.isnan(train_loss), 'Model diverged with loss = NaN' if step % steps_per_epoch == 0 or (step + 1) == g_config['max_steps']: checkpoint_path = os.path.join(g_config['train_dir'], 'model.ckpt') saver.save(sess, checkpoint_path, global_step=step)
def ckpt_pb(pb_path, lite_path): config = tf.ConfigProto(allow_soft_placement=True) config.gpu_options.allow_growth = True # to pb tf.reset_default_graph() with tf.Graph().as_default() as graph, tf.Session(graph=graph) as sess: path_base = Path(g_config['eval_dataset']).parent.parent _mean_shape = mio.import_pickle(path_base / 'reference_shape.pkl') _mean_shape = data_provider.align_reference_shape_to_112(_mean_shape) assert isinstance(_mean_shape, np.ndarray) print(_mean_shape.shape) tf_img = tf.placeholder(dtype=tf.float32, shape=(1, 112, 112, 3), name='Inputs/InputImage') tf_dummy = tf.placeholder(dtype=tf.float32, shape=(1, 73, 2), name='dummy') tf_shape = tf.constant(_mean_shape, dtype=tf.float32, shape=(73, 2), name='Inputs/MeanShape') mdm_model.MDMModel( tf_img, tf_dummy, tf_shape, batch_size=1, num_patches=g_config['num_patches'], num_channels=3, is_training=False ) saver = tf.train.Saver() ckpt = tf.train.get_checkpoint_state(g_config['train_dir']) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) global_step = ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1] print('Successfully loaded model from {} at step={}.'.format(ckpt.model_checkpoint_path, global_step)) else: print('No checkpoint file found') return output_graph_def = tf.graph_util.convert_variables_to_constants( sess, tf.get_default_graph().as_graph_def(), ['Network/Predict/add'] ) with tf.gfile.FastGFile(pb_path, mode='wb') as f: f.write(output_graph_def.SerializeToString()) # to tflite tf.reset_default_graph() with tf.Graph().as_default() as graph: with tf.gfile.GFile(pb_path, "rb") as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) tf.import_graph_def( graph_def, input_map=None, return_elements=None, name='prefix', op_dict=None, producer_op_list=None ) tf.summary.FileWriter('.', graph=graph) converter = tf.contrib.lite.TFLiteConverter.from_frozen_graph( pb_path, ['Inputs/InputImage'], ['Network/Predict/Reshape', 'Inputs/MeanShape'] ) with tf.gfile.FastGFile(lite_path, mode='wb') as f: f.write(converter.convert()) # Load TFLite model and allocate tensors. interpreter = tf.contrib.lite.Interpreter(model_path=lite_path) interpreter.allocate_tensors() # Get input and output tensors. input_details = interpreter.get_input_details() output_details = interpreter.get_output_details() print(input_details) print(output_details)