Auto Mesh TensorFlow - Vengineerの戯言

@Vengineerの戯言 : Twitter
SystemVerilogの世界へようこそ、すべては、SystemC v0.9公開から始まった

Mesh TensorFlow でも AutoML なる Auto Mesh TensorFlow もあるね。

コードは、こんな感じ。。。

import mesh_tensorflow.auto_mtf

graph = mtf.Graph()
mesh = mtf.Mesh(graph, "my_mesh")
# Insert model code here.
outputs = [logits, loss] # iterable of mtf.Tensor, the outputs you're computing
mesh_shape = [("processor_rows", 2), ("processor_cols", 2)]
layout_rules = mtf.auto_mtf.layout(graph, mesh_shape, outputs)

mtf.auto.mtf.layout がそれっぽい。。。

・Auto MTF currently tries to choose the layout which minimizes the peak memory usage of the computation

とあるので、現時点ではピーク時のメモリを最小にするために頑張るのね。

・In the future, we hope to also support layouts which result in efficient computations while keeping peak memory under a threshold.

ともあるけど。

layout メソッドの他に、layout_and_mesh_shape メソッドもあるようだ。

実際には、LayoutOptimizer クラスが頑張ってくれるみたいだ。