8000 minigo/dual_net.py at master · zihevic/minigo · GitHub
[go: up one dir, main page]
Skip to content

dual_net.py

Latest commit

 

History

History
774 lines (634 loc) · 29 KB

dual_net.py

File metadata and controls

774 lines (634 loc) · 29 KB
Raw
def get_estimator():
if FLAGS.use_tpu:
return _get_tpu_estimator()
else:
return _get_nontpu_estimator()
def _get_nontpu_estimator():
session_config = tf.ConfigProto()
session_config.gpu_options.allow_growth = True
run_config = tf.estimator.RunConfig(
save_summary_steps=FLAGS.summary_steps,
keep_checkpoint_max=FLAGS.keep_checkpoint_max,
session_config=session_config)
return tf.estimator.Estimator(
model_fn,
model_dir=FLAGS.work_dir,
config=run_config,
params=FLAGS.flag_values_dict())
def _get_tpu_estimator():
tpu_cluster_resolver = contrib_cluster_resolver.TPUClusterResolver(
FLAGS.tpu_name, zone=None, project=None)
tpu_grpc_url = tpu_cluster_resolver.get_master()
run_config = contrib_tpu_python_tpu_tpu_config.RunConfig(
master=tpu_grpc_url,
evaluation_master=tpu_grpc_url,
model_dir=FLAGS.work_dir,
save_checkpoints_steps=max(1000, FLAGS.iterations_per_loop),
save_summary_steps=FLAGS.summary_steps,
keep_checkpoint_max=FLAGS.keep_checkpoint_max,
session_config=tf.ConfigProto(
allow_soft_placement=True, log_device_placement=True),
tpu_config=contrib_tpu_python_tpu_tpu_config.TPUConfig(
iterations_per_loop=FLAGS.iterations_per_loop,
num_shards=FLAGS.num_tpu_cores,
per_host_input_for_training=contrib_tpu_python_tpu_tpu_config.InputPipelineConfig.PER_HOST_V2))
return contrib_tpu_python_tpu_tpu_estimator.TPUEstimator(
use_tpu=FLAGS.use_tpu,
model_fn=model_fn,
config=run_config,
train_batch_size=FLAGS.train_batch_size * FLAGS.num_tpu_cores,
eval_batch_size=FLAGS.train_batch_size * FLAGS.num_tpu_cores,
params=FLAGS.flag_values_dict())
def bootstrap():
"""Initialize a tf.Estimator run with random initial weights."""
# a bit hacky - forge an initial checkpoint with the name that subsequent
# Estimator runs will expect to find.
#
# Estimator will do this automatically when you call train(), but calling
# train() requires data, and I didn't feel like creating training data in
# order to run the full train pipeline for 1 step.
maybe_set_seed()
initial_checkpoint_name = 'model.ckpt-1'
save_file = os.path.join(FLAGS.work_dir, initial_checkpoint_name)
sess = tf.Session(graph=tf.Graph())
with sess.graph.as_default():
features, labels = get_inference_input()
model_fn(features, labels, tf.estimator.ModeKeys.PREDICT,
params=FLAGS.flag_values_dict())
sess.run(tf.global_variables_initializer())
tf.train.Saver().save(sess, save_file)
def export_model(model_path):
"""Take the latest checkpoint and copy it to model_path.
Assumes that all relevant model files are prefixed by the same name.
(For example, foo.index, foo.meta and foo.data-00000-of-00001).
Args:
model_path: The path (can be a gs:// path) to export model
"""
estimator = tf.estimator.Estimator(model_fn, model_dir=FLAGS.work_dir,
params=FLAGS.flag_values_dict())
latest_checkpoint = estimator.latest_checkpoint()
all_checkpoint_files = tf.gfile.Glob(latest_checkpoint + '*')
for filename in all_checkpoint_files:
suffix = filename.partition(latest_checkpoint)[2]
destination_path = model_path + suffix
print('Copying {} to {}'.format(filename, destination_path))
tf.gfile.Copy(filename, destination_path)
def freeze_graph(model_path, use_trt=False, trt_max_batch_size=8,
trt_precision='fp32'):
output_names = ['policy_output', 'value_output']
n = DualNetwork(model_path)
out_graph = tf.graph_util.convert_variables_to_constants(
n.sess, n.sess.graph.as_graph_def(), output_names)
if use_trt:
import tensorflow.contrib.tensorrt as trt
out_graph = trt.create_inference_graph(
input_graph_def=out_graph,
outputs=output_names,
max_batch_size=trt_max_batch_size,
max_workspace_size_bytes=1 << 29,
precision_mode=trt_precision)
metadata = make_model_metadata({
'engine': 'tf',
'use_trt': bool(use_trt),
})
minigo_model.write_graph_def(out_graph, metadata, model_path + '.minigo')
def freeze_graph_tpu(model_path):
"""Custom freeze_graph implementation for Cloud TPU."""
assert model_path
assert FLAGS.tpu_name
if FLAGS.tpu_name.startswith('grpc://'):
tpu_grpc_url = FLAGS.tpu_name
else:
tpu_cluster_resolver = contrib_cluster_resolver.TPUClusterResolver(
FLAGS.tpu_name, zone=None, project=None)
tpu_grpc_url = tpu_cluster_resolver.get_master()
sess = tf.Session(tpu_grpc_url)
output_names = []
with sess.graph.as_default():
# Replicate the inference function for each TPU core.
replicated_features = []
feature_type = tf.bool if FLAGS.bool_features else tf.float32
for i in range(FLAGS.num_tpu_cores):
name = 'pos_tensor_%d' % i
features = tf.placeholder(
feature_type, [None], name=name)
replicated_features.append((features,))
outputs = contrib_tpu.replicate(
tpu_model_inference_fn, replicated_features)
# The replicate op assigns names like output_0_shard_0 to the output
# names. Give them human readable names.
for i, (policy_output, value_output, _) in enumerate(outputs):
policy_name = 'policy_output_%d' % i
value_name = 'value_output_%d' % i
output_names.extend([policy_name, value_name])
tf.identity(policy_output, policy_name)
tf.identity(value_output, value_name)
tf.train.Saver().restore(sess, model_path)
out_graph = tf.graph_util.convert_variables_to_constants(
sess, sess.graph.as_graph_def(), output_names)
metadata = make_model_metadata({
'engine': 'tpu',
'num_replicas': FLAGS.num_tpu_cores,
})
minigo_model.write_graph_def(out_graph, metadata, model_path + '.minigo')
def make_model_metadata(metadata):
for f in ['conv_width', 'fc_width', 'trunk_layers', 'use_SE', 'use_SE_bias',
'use_swish', 'input_features', 'input_layout']:
metadata[f] = getattr(FLAGS, f)
metadata['input_type'] = 'bool' if FLAGS.bool_features else 'float'
metadata['board_size'] = go.N
return metadata
0