Distributed TensorFlow TensorFlow User Group Tokyo #8 @shuhei_fujiwara 2018-09-14 1

自己紹介 藤原秀平 (FUJIWARA Shuhei) Twitter: @shuhei_fujiwara ▶ 株式会社メルペイ Machine Learning Team ▶ TensorFlow User Group Tokyo Organizer ▶ Google Developer Expert (Machine Learning) 2

機械学習でよくある並列化 正則化項 + 損失の最小化問題: min x λR(w) + L(w) where L(w) = ∑ i ℓi(w) 損失項の勾配はデータごとに並列で計算できる wk+1 = wk − α [ ∇wR(wk) + ∑ i ∇wℓi(wk) ] 3

分散学習の役割分担の例 Parameter Server ▶ 最新の weight w を保持して必要なときに master や worker に共有 Worker ▶ それぞれ異なるデータを使って勾配を計算 Master ▶ Worker とだいたい同じ ▶ モデルの保存や初期化など実装上必要な一部の処理 4

データ並列な同期的分散学習 Compute grad. grad. grad. sum Share Update Master and Workers Parameter Servers mini-batch samples 5

データ並列な非同期的分散学習 Compute grad. grad. grad. Share Update Master and Workers Parameter Servers mini-batch samples 6

それぞれの長所 非同期型 ▶ 対故障性 ▶ Worker が落ちても影響が小さく学習が継続可能 ▶ 待ち合わせが無いので高スループット ▶ 遅い worker に足を引っ張られない 同期型 ▶ 収束性が良く学習が上手くいきやすい ▶ Stale gradients の問題が発生しない 7

TensorFlow で実装するときのお作法 Cluster Worker Worker PS PS Master task.py task.py task.py task.py task.py ▶ クラスタの全ノードは同じスクリプトを実行 ▶ スクリプト内で master, worker, ps の分岐を記述 ▶ 自分がどれなのかという情報は環境変数に入れておく 8

環境変数に入れておく情報の例 1 { 2 "cluster": { 3 "ps": ["ps-408715a125-0:2222", "ps-408715a125-1:2222"], 4 "worker": ["worker-408715a125-0:2222", "worker-408715a125-1:2222"], 5 "master": ["master-408715a125-0:2222"] 6 }, 7 "task": { 8 "index": 0, 9 "type": "master", 10 ... 11 }, 12 ... 13 } 9

TensorFlow 側の処理 1 tf_conf = json.loads(os.environ.get("TF_CONFIG", "{}")) 2 # TF_CONFか ら ク ラ ス タ 構 成 の 情 報 を 取 り 出 し て 使 う 3 server = tf.train.Server( 4 tf_conf.get("cluster", None), job_name=tf_conf["task"]["type"], 5 task_index=tf_conf["task"]["index"] 6 ) 7 # TF_CONFか ら 自 分 の 役 割 がMaster、Worker、PSの ど れ か 調 べ て 分 岐 さ せ る 8 if tf_conf["task"]["type"] == "ps": 9 server.join() # Parameter Server は こ の 処 理 を す る だ け 10 else: 11 # MasterとWorkerの 処 理 を 実 行 す る 12 ... 10

分散学習の実装は楽なのか？ ▶ なんだかんだ言ってもどっかで嵌まる ▶ 環境を用意するのが面倒くさい ▶ クラスタの構築 ▶ GPU を使った計算環境 ▶ 高速なネットワークを用意できないと台無し 11

Estimator ▶ TensorFlow の高レベル API ▶ コードを変更せずに分散学習できる 1 clf = tf.estimator.DNNClassifier( 2 feature_columns=[tf.feature_column.numeric_column("x", shape=[4]), 3 hidden_units=[10, 20, 10], 4 n_classes=3, 5 model_dir="/tmp/iris_model" 6 ) 7 clf.train(...) 12

Google Cloud Machine Learning Engine ▶ 環境変数とか全部入りのクラスタがにゅっと生えてくる ▶ TPU も最近使えるようになった ▶ Preemptible VM も最近使えるようになった ▶ 最近何か安くなった 13

All Reduce 同期型に絞るなら別のクラスタ構成もある ▶ Weight w は各 worker が コピーを持つ ▶ Worker 同士で通信して 勾配の情報を交換 ▶ それぞれ w を更新する ▶ Chainer MN はこの方式 gradient with all samples mini-batch samples 14

TensorFlow で All Reduce 最近追加された！ ▶ Distribution strategy を Estimator に渡すだけ ▶ Multi node 対応はまだなので 1 台に複数 GPU という状況を想定 1 distribution = tf.contrib.distribute.MirroredStrategy(num_gpus=2) 2 run_config = tf.estimator.RunConfig(distribute=distribution) 3 classifier = tf.estimator.Estimator( 4 model_fn=model_function, 5 model_dir=model_dir, 6 config=run_config 7 ) 15

まとめ ▶ TensorFlow を使うと分散学習のコードが（そこそこ）簡単に書ける ▶ 最近は Estimator がけっこう良い感じ ▶ ML Engine を使うと分散学習用のクラスタを用意するのが楽 16