我正在研究 XLA,并希望直观地了解 XLA 执行的各种优化,特别是内核融合。我一直在使用以下标志转储图形并将它们导入 tensorboard,但这些似乎是优化前的图形。我确认 XLA 正在融合来自同样被转储的 HLA IR 的内核。
我正在使用的环境变量及其值
TF_DUMP_GRAPH_PREFIX=hlo TF_XLA_FLAGS="--tf_xla_clustering_debug --tf_xla_auto_jit=2"XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=./hlo
最佳答案
将以下代码片段添加到您要检查的优化过程中。 DumpHloModuleIfEnabled
会将 HLO 图发送到 ./hlo
,对于任何 hlo 模块。
DumpHloModuleIfEnabled(*module, "your_hlo_module_to_print");
// use case:
StatusOr<bool> CustomGraphRewriter::Run(HloModule* module) {
// this will dump the graph before the optimization
DumpHloModuleIfEnabled(*module, "before_graph_rewriting");
bool is_rewriting_done = false;
TF_ASSIGN_OR_RETURN(is_rewriting_done,
RunGraphRewriting(module));
// this will dump the graph after the optimization
DumpHloModuleIfEnabled(*module, "after_graph_rewriting");
return is_rewriting_done;
}
对于可视化工具,使用 graphviz 会有所帮助。安装 graphviz 后,您应该将图形发送到 *.dot 文件,并带有以下标志。
XLA_FLAGS="--xla_dump_hlo_as_dot"
为了可视化点文件,这会起作用:
dot -Tsvg your_hlo_graph.dot > visualized_output.svg
在那里,这将发出如图所示的结果。 sample HLO graph visualized
关于tensorflow - 可视化优化前后的 XLA 图,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/58650149/