在观看了TF Dev Summit 演讲，了解到在 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)

然而，文档警告说这种方法仅用于测试。我希望能够使用推荐的方式来实现这一点：

# 配置以启用 JIT 编译config = tf.ConfigProto()config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1sess = tf.Session(config=config)

但我无法找到在 model_fn 外部设置 tf.Session 的方法。

以下伪代码可能更能说明问题：

def char_cnn_model(features, target, mode, params, model_dir):    # 构建 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):    # 加载训练数据、测试数据等    ...    classifier = learn.Estimator(        model_fn=char_cnn_model,  # 如上所定义        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 在 main 中定义    accuracy_score = classifier.evaluate(x=X_test, y=y_test)                                      

tf.contrib.learn.RunConfig 看起来是一个不错的选择，但它没有关于会话的设置（这似乎是合理的，为什么 RunConfig 会有一个指向现有会话的指针呢？）

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

回答：

看起来目前还不行。MonitoredSession 是直接实例化的，选项是本地传递的。你唯一的选择是子类化并用一个使用 XLA 配置的监控会话覆盖 train 方法。