TensorFlowをKubernetesで分散学習するためのKubeFlowを動かしてみた話

Kubernetesが機械学習基盤としてKubeflowをリリースしました。

ざっくりいうと、KubeFlowは以下の機械学習アプリケーション開発の一連のワークロードをサポートするOSSです。

• モデル開発基盤

データサイエンティストや機械学習エンジニアがJupyterNotebookをつかってモデル開発するための、Jupyterが動くサーバ環境を提供

• TensorFlow分散学習基盤

開発したモデルをTensorFlowで分散学習するためのKubernetesクラスタを自動生成

• アプリ公開基盤

学習済みモデルをKubernetesクラスタ上でサービス公開する基盤を提供(Tensorflow Serving)

github.com

まだ、開発途中、かつざっとしか追えていませんが、導入のしかたと概要をまとめます。

0. kubeFlow環境構築

Kubeflowを動かすためには、まずKubernetesクラスタを構築しなければいけません。 今回は検証のためなので、雑にminikubeを使っています。

手順は以下のとおりです。

minikubeのインストール

MacOSのときは、次のコマンドでminikubeをインストールします。

$ brew cask install minikube

Kubernetesクラスタの構築

次のコマンドを実行して、Kubernetesクラスタを構築します。

$ minikube start
Starting local Kubernetes v1.8.0 cluster...
Starting VM...
Getting VM IP address...
Moving files into cluster...
Setting up certs...
Connecting to cluster...
Setting up kubeconfig...
Starting cluster components...
Kubectl is now configured to use the cluster.
Loading cached images from config file.

バージョンは次の通りです。

$ minikube version
minikube version: v0.24.1

$ kubectl get node
NAME       STATUS    AGE       VERSION
minikube   Ready     10m       v1.8.0

これで準備は完了です。

なお、基本的人権を満たさないPCを使っているみなさまと プロキシ環境で無駄に人生を消耗しているみなさまは AWSやAzureやGCPなどのパブリッククラウドのVMかKubernetesマネージドサービス使ってください。

Kubeflowのインストール

KubeflowはGitHubで公開されていますので、次のコマンドでクローンします。

$ git clone https://github.com/google/kubeflow

ここでは、CPUを使って分散学習する例を説明します。 次のコマンドを実行して、環境を構築します。

$ cd kubeflow
$ kubectl apply -f components/ -R

なお、KubeflowはGPUもサポートしています。正確には、Kubernetes1.8からGPU対応がα提供されているので、Nvidia CUDAなどで後日動かしてみたいとおもいます。

環境構築のためのマニュフェストファイルは、kubeflow/components配下にあります。

1. モデル開発基盤

KubeFlowではモデル開発のため、JupyterNotebookが動作する環境が用意されます。 KubeFlowは、JupyterHubを使用しています。なので複数のユーザーでの認証アクセスを管理でき、「spawners」と呼ばれるプラグイン可能なコンポーネントをつかってます。

JupyterNotebookはKubernetesのPodで動きます。

$ kubectl get pod
NAME                                                READY     STATUS    RESTARTS   AGE
~中略~
tf-hub-0                                            1/1       Running   0          2h
tf-job-operator-5c648c58bc-vrtcx                    1/1       Running   0          2h

JupyterNotebookにアクセスするためのPodは「tf-hub-lb」という名前のServiceがLoadBaranceとして公開されています。

tf-cnn $ kubectl get svc
NAME                                          CLUSTER-IP       EXTERNAL-IP   PORT(S)          AGE
～中略～
tf-hub-0                                      None             <none>        8000/TCP         2h
tf-hub-lb                                     10.107.227.23    <pending>     80:32089/TCP     2h

minikubeの場合は、次のコマンドで公開URLを確認できます。

$ minikube service tf-hub-lb --url
http://192.168.99.100:32089

カスタムのJupyterNotebookを使いたいときや構成を変更したいときは、以下のマニュフェストを修正します。 kubeflow/components/jupyterhub/manifests/

2. TensorFlow分散学習基盤

チュートリアルや書籍のサンプル程度のモデルであれば、 使っているPCやクラウドのVMのシングルインスタンスで大丈夫ですが 実務で使うような大規模なものであれば、計算量も膨大になるため、(コストを見ながら)計算基盤をスケールさせる必要があります。

TensorFlowの分散学習処理を行う際は、3種類のノードをつかいます。

• Master
• Parameter Server
• Worker

Kubeflowでは、Kubernetesクラスタをつかって分散クラスタを構成し、その上でTensorFlowの分散処理コードを動かす環境が用意されています。

MasterとParameter Serverは1〜2ノード、Workerは計算に必要な沢山のノードが動きます。

Kubeflowのサンプルコードが用意されているので次のコマンドで実行します。 このサンプルは、tf_cnn_benchmarksを使用して畳み込みニューラルネットワークを学習するためのものです。

ソースはこちらですね。

benchmarks/scripts/tf_cnn_benchmarks at master · tensorflow/benchmarks · GitHub

次のコマンドでKubernetesクラスタにJobを投入します。

$ cd kubeflow/tf-controller-examples/tf-cnn/
$ kubectl create -f tf_job_cpu_distributed.yaml 
tfjob "inception-171202-163257-cpu-3" created

しばらくすると、次のようなPodが生成されます。

kubectlコマンドで確認すると、inception--master-/inception--ps-がそれぞれ1つずつ inception--worker-が3つできているのがわかります。

$ kubectl get pod
NAME                                                READY     STATUS              RESTARTS   AGE
inception-171202-163257-cpu-3-master-vjo4-0-hc2kl   0/1       ContainerCreating   0          18s
inception-171202-163257-cpu-3-ps-vjo4-0-728vm       0/1       ContainerCreating   0          18s
inception-171202-163257-cpu-3-worker-vjo4-0-h66bn   0/1       ContainerCreating   0          18s
inception-171202-163257-cpu-3-worker-vjo4-1-lcljx   0/1       ContainerCreating   0          18s
inception-171202-163257-cpu-3-worker-vjo4-2-j59bh   0/1       ContainerCreating   0          18s
model-server-6598c6486d-b8rxg                       1/1       Running             1          6m
model-server-6598c6486d-clwz9                       0/1       Pending             0          6m
model-server-6598c6486d-dkmsk                       0/1       Pending             0          6m
tf-hub-0                                            1/1       Running             1          6m
tf-job-operator-5c648c58bc-47c7f                    1/1       Running             1     

KubeFlowでつかうマニフェストファイルは以下のようになっています。 ポイントとなるところだけ抜粋します。

apiVersion

apiVersionで「tensorflow.org/v1alpha1」を、kindで「TFJob」を指定してます。 kindは「Job」ではないので注意です。

apiVersion: tensorflow.org/v1alpha1
kind: TfJob

Master Pod

MasterのPodをreplicaSpecsで指定します。 実行したいジョブのコマンドやパラメータの指定で、通常のKubernetesマニュフェストファイルと同じですが [tfReplicaType]を「MASTER」にしています。 また、[replicas]が1なので、Podが1つ上がります。

spec:
  replicaSpecs:
  - replicas: 1
    template:
      spec:
        containers:
        - args:
          - python
          - tf_cnn_benchmarks.py
          - --batch_size=32
          - --model=resnet50
          - --variable_update=parameter_server
          - --flush_stdout=true
          - --num_gpus=1
          - --local_parameter_device=cpu
          - --device=cpu
          - --data_format=NHWC
          image: gcr.io/kubeflow/tf-benchmarks-cpu:v20171202-bdab599-dirty-284af3
          name: tensorflow
          workingDir: /opt/tf-benchmarks/scripts/tf_cnn_benchmarks
        restartPolicy: OnFailure
    tfReplicaType: MASTER

Parameter Server

同様にParameter ServerのPodを定義しています。 Parameter Serverは[tfReplicaType]が「PS」で [replicas]が1なので、Master同様Podが1つ上がります。

  - replicas: 1
    template:
      spec:
        containers:
        - args:
          ~中略~
          name: tensorflow
          workingDir: /opt/tf-benchmarks/scripts/tf_cnn_benchmarks
        restartPolicy: OnFailure
    tfReplicaType: PS
  tfImage: gcr.io/kubeflow/tf-benchmarks-cpu:v20171202-bdab599-dirty-284af3

Worker Pod

最後に計算Podです。 Parameter Serverは[tfReplicaType]が「WORKER」で [replicas]が3なので、計算用Podが3つ上がります。

  - replicas: 3
    template:
      spec:
        containers:
        - args:
          ~中略~
          name: tensorflow
          workingDir: /opt/tf-benchmarks/scripts/tf_cnn_benchmarks
        restartPolicy: OnFailure
    tfReplicaType: WORKER

この構成で分散して学習が行われます。

今回はKubernetesクラスタをminikubeで構成していますが、たとえばGKEなどをつかうときは、NodePoolを構成して、プール内に適切なスペックのノードをおけばよいかとおもいます。

TensorFlowの分散学習の実装については、enakaiさんのブログがわかりやすいです。

分散学習用TensorFlowコードの書き方 - めもめも

3. アプリ公開基盤

KubeFlowのアプリ公開には、Tensorflow Servingが使われています。 今回は、minikubeで確認したので試してませんが GCPなどでクラスタを構成して 学習済みのモデルをGCSにアップロードしてアプリで利用する流れです。

github.com

DevFest2017で、「プログラマのためのGoogle Cloud Platform超入門」というタイトルで GKEを簡単に紹介するセッションをさせていただきましたが そこで、TensorFlowで画像推論する基盤をデモとして動かしました。

ほぼ、同じようなアーキテクチャなのでご参考までに。

感想

世間やメディアで大変な賑わいをみせている深層学習界隈は、人間がAIに仕事を奪われるたぐいのファンタジーからはじまり、 クールなアルゴリズムの話題、論文を実装してみました速報、ツール/フレームワークを使ってみました日記や 深層学習でこんなすごいことできました選手権などの情報にあふれ、とてもきらきらまぶしく膨大過ぎて追いかけきれてないのですが、、、、

このような基盤構築や地味な方法論は、とても楽しいので、細々と調べながら手を動かしたいと思いました。

まだ開発途中のようですが、ひそかにウォッチしたいとおもいます。

おわり