python - 使用 tf.contrib.learn.Estimator 时启用 XLA

Question

看到 TF Dev Summit talk关于在 TensorFlow 图表上启用 XLA 的潜在(和高度实验性)好处，我想我会稍微试验一下。

问题:使用tf.contrib.learn.Estimator时, 如何启用 JIT XLA ？

我可以通过

为 JIT XLA 标记某些操作

with tf.device("/job:localhost/replica:0/task:0/device:XLA_GPU:0"):
    output = tf.add(input1, input2)

但是，文档警告说这只是为了测试。我希望能够使用推荐的方式执行此操作

# Config to turn on JIT compilation
config = tf.ConfigProto()
config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1

sess = tf.Session(config=config)

但我想不出从 model_fn 外部设置 tf.Session 的方法。

这段伪代码可以更好地阐明问题:

def char_cnn_model(features, target, mode, params, model_dir):
    # do stuff to build the CNN

    ...

    return tf.contrib.learn.ModelFnOps(mode=mode,
        predictions=predictions_dict,
        loss=loss,
        train_op=train_op,
        eval_metric_ops=eval_metric_ops)

def main(unused_argv):
    # load training data, test data etc

    ...

    classifier = learn.Estimator(
        model_fn=char_cnn_model,  # defined above
        model_dir=model_dir,
        params=params,
        config=tf.contrib.learn.RunConfig(save_checkpoints_secs=60,
                                          log_device_placement=True,
                                          tf_random_seed=7))
    classifier.fit(
        X_train,
        y_train,
        steps=5000,
        monitors=[validation_monitor])  # validation_monitor defined in main

    accuracy_score = classifier.evaluate(x=X_test, y=y_test)                                      

tf.contrib.learn.RunConfig似乎是一个很好的候选者，但它没有用于 session 的东西(我想这是有道理的，为什么 RunConfig 会有一个指向现有 session 的指针？)

我可能想得太多了 tf.get_default_session可能是我所需要的，但我可以在 session 创建后修改它的配置吗？

Answer 1

现在可以通过 RunConfig 的 session_config 参数实现.这是一个例子:

session_config = tf.ConfigProto()                                               
optimizer_options = session_config.graph_options.optimizer_options               
if xla:                                                     
    optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1           

run_config = tf.estimator.RunConfig(                                        
    session_config=session_config,  # use this session config
    log_device_placement=True                                           
)                                                                           

base_classifier = tf.estimator.Estimator(                                   
    model_fn=model_fn,                                                      
    model_dir=model_dir,                                                    
    config=run_config,                                                      
    params=model_params                                                                
)