AWS Batch環境でのAPV-MCTSのセルフプレイによる強化学習をGPU化してみた
前回の記事で作成したAPV-MCTSのセルフプレイによる強化学習のAWS Batch環境をGPU化して動かしてみました。
今回のコード
下記、タグv0.0.2になります。
環境の概要
全体構成は下記です。
前回の構成から下記の変更を行いました。
- nvidia-docker2に対応したAMIの作成を行うCodeBuildを追加
- セルフプレイの新モデル作成時のトレーニングをGPUに対応
- AWS Batchのコンピューティング環境で使うAMIとインスタンスを変更
GPU化の実装
具体的な実装内容です。ちなみに環境作成手順はbashからマークダウン形式に変更しdocs/setup_batch.mdに移動しました。
nvidia-docker2に対応したAMIの作成を行うCodeBuildを追加
Packerを使ってDeep Learning AMI(ami-08a7740ff4d3fd90f)をベースにnvidia-docker2に対応したAMIを作成するCodeBuildを追加しました。
CodeBuildでPackerを使ってのAMIの作成方法は下記を参考にしました。
nvidia-docker2に対応したAMIの作成方法は下記を参考にしました。
ただ、Deep Learning AMIの/usr/local/cuda/version.txtが複数行になっていたのでconfigure-gpu.shのGet CUDA versionはcatをhead -n 1に変更しました。
以下、コードです。
CodeBuildの作成手順は下記です。Packerで必要になるAMI作成場所のサブネットID*1などを環境変数で渡すように指定しています。
cat <<EOF > Source.json
{
"type": "CODECOMMIT",
"location": "https://git-codecommit.ap-northeast-1.amazonaws.com/v1/repos/heta-reversi",
"buildspec": "buildspec_ami.yml"
}
EOF
cat <<EOF > Artifacts.json
{
"type": "NO_ARTIFACTS"
}
EOF
cat <<EOF > Environment.json
{
"type": "LINUX_CONTAINER",
"image": "aws/codebuild/docker:18.09.0",
"computeType": "BUILD_GENERAL1_SMALL",
"environmentVariables": [
{
"name": "AWS_REGION",
"value": "ap-northeast-1",
"type": "PLAINTEXT"
},
{
"name": "AWS_SUBNET_ID",
"value": "${SUBNET_ID}",
"type": "PLAINTEXT"
}
]
}
EOF
aws codebuild create-project --name test-batch-ami \
--source file://Source.json \
--artifacts file://Artifacts.json \
--environment file://Environment.json \
--service-role ${ROLE_AMI_BUILD_ARN}
CODEBUILD_ID=$(aws codebuild start-build --project-name test-batch-ami --source-version ${BRANCH} | tr -d "\n" | jq -r '.build.id')
echo "started.. id is ${CODEBUILD_ID}"
while true
do
sleep 10s
STATUS=$(aws codebuild batch-get-builds --ids "${CODEBUILD_ID}" | tr -d "\n" | jq -r '.builds[].buildStatus')
echo "..status is ${STATUS}."
if [ "${STATUS}" != "IN_PROGRESS" ]; then
if [ "${STATUS}" != "SUCCEEDED" ]; then
echo "faild."
fi
echo "done."
break
fi
done
AMI_IDS=$(aws ec2 describe-images \
--owner self \
--filters "Name=tag:USE,Values=heta-reversi" | jq -r ".Images[].ImageId")
AMI_ID="${AMI_IDS[0]}"
CodeBuildで実施するAMI作成処理は下記です。今回はAMIが１つしかない前提でのAMI ID取得*2になっているためAMI作成前に過去に作成済みのAMIを削除するようにしてあります。*3
version: 0.2 phases: pre_build: commands: - curl -qL -o packer.zip https://releases.hashicorp.com/packer/1.3.3/packer_1.3.3_linux_amd64.zip && unzip packer.zip - ./packer validate amazon-linux_packer-template.json build: commands: - ./delete_amis.sh - ./packer build amazon-linux_packer-template.json post_build: commands: - echo "done $(date)"
過去のAMIを削除する処理は下記です。今回の作成AMIにはタグUSEにheta-reversiを付与するようにしてあるため削除の目印にしてあります。
#!/bin/bash
aws ec2 describe-images \
--owner self \
--filters "Name=tag:USE,Values=heta-reversi" | jq -r ".Images[].ImageId" | while read -r AMI_ID
do
aws ec2 deregister-image --image-id ${AMI_ID}
done
AMIを作成するPackerのTemplateは下記です。インスタンスタイプはGPUインスタンスのp2.xlargeを使うようにしました。
amazon-linux_packer-template.json
{ "variables": { "aws_region": "{{env `AWS_REGION`}}", "aws_subnet_id": "{{env `AWS_SUBNET_ID`}}", "aws_ami_name": "amazon-linux_heta-reversi" }, "builders": [{ "type": "amazon-ebs", "region": "{{user `aws_region`}}", "instance_type": "p2.xlarge", "subnet_id": "{{user `aws_subnet_id`}}", "ssh_username": "ec2-user", "ami_name": "{{user `aws_ami_name`}}", "ami_description": "heta-reversi's Amazon Linux", "source_ami": "ami-08a7740ff4d3fd90f", "ssh_pty": true, "tags": { "USE": "heta-reversi" } }], "provisioners": [ { "type": "file", "source": "configure-gpu.sh", "destination": "/home/ec2-user/configure-gpu.sh" }, { "type": "shell", "inline": [ "bash ./configure-gpu.sh", "sudo rm -rf /var/lib/ecs/data/ecs_agent_data.json" ] } ] }
PackerのTemplateのprovisionersで実行するnvidia-docker2の設定処理は下記です。dockerのdefault-runtimeにnvidiaを指定してnvidia-docker2が有効になるようにしています。
#!/bin/bash
# Install ecs-init, start docker, and install nvidia-docker 2
sudo yum install -y ecs-init
sudo service docker start
DOCKER_VERSION=$(docker -v | awk '{ print $3 }' | cut -f1 -d"-")
DISTRIBUTION=$(. /etc/os-release;echo ${ID$VERSION_ID})
curl -s -L https://nvidia.github.io/nvidia-docker/${DISTRIBUTION}/nvidia-docker.repo | \
sudo tee /etc/yum.repos.d/nvidia-docker.repo
PACKAGES=$(sudo yum search -y --showduplicates nvidia-docker2 nvidia-container-runtime | grep ${DOCKER_VERSION} | awk '{ print $1 }')
sudo yum install -y ${PACKAGES}
sudo pkill -SIGHUP dockerd
# Get CUDA version
CUDA_VERSION=$(head -n 1 /usr/local/cuda/version.txt | awk '{ print $3 }' | cut -f1-2 -d".")
# Run test container to verify installation
sudo docker run --privileged --runtime=nvidia --rm nvidia/cuda:${CUDA_VERSION}-base nvidia-smi
# Update Docker daemon.json to user nvidia-container-runtime by default
sudo tee /etc/docker/daemon.json <<EOF
{
"runtimes": {
"nvidia": {
"path": "/usr/bin/nvidia-container-runtime",
"runtimeArgs": []
}
},
"default-runtime": "nvidia"
}
EOF
sudo service docker restart
セルフプレイの新モデル作成時のトレーニングをGPUに対応
現在のモデルをベースにセルフプレイで生成した棋譜データを使って学習している部分をGPUで行うようにしました。変更を行なったらCodeBuildを実行してECRのdocker imageを更新します。
以下、コードです。
まずGPUを使う箇所を局所化するためにgpu_deviceをグローバル定義しました。*4この値を目印にGPU処理にするか判断します。
gpu_device = -1
学習を行う_create_new_modelをGPU化します。学習の実施前後でモデルをto_gpuでGPUモードにしてto_cpuでCPUモードに戻しています。
def _create_new_model(self, steps_list, epoch_num = None, batch_size = 2048): epoch_num = default_params['dual_net_trainer_create_new_model_epoch_num'] if epoch_num is None else epoch_num model = DualNet() model.load(self.model_filename) if gpu_device >= 0: cuda.get_device(gpu_device).use() model.to_gpu(gpu_device) optimizer = chainer.optimizers.Adam() optimizer.setup(model) for i in range(epoch_num): x_train, y_train_policy, y_train_value = self._get_train_batch(steps_list, batch_size) y_policy, y_value = model(x_train) model.cleargrads() loss = F.mean_squared_error(y_policy, y_train_policy) + F.mean_squared_error(y_value, y_train_value) loss.backward() optimizer.update() print("[new nodel] epoch: {} / {}, loss: {}".format(i + 1, epoch_num, loss)) if gpu_device >= 0: model.to_cpu() return model
棋譜データから学習データを取り出す_get_train_batchです。GPUモードの場合はnumpyでは無くcuda.cupyを使う必要があります。
def _get_train_batch(self, steps_list, batch_size): batch_x, batch_y_policy, batch_y_value = [], [], [] for _ in range(batch_size): x, y_policy, y_value = self._get_train_random(steps_list) batch_x.append(x) batch_y_policy.append(y_policy) batch_y_value.append(y_value) xp = np if gpu_device >= 0: xp = cuda.cupy x_train = Variable(xp.array(batch_x)).reshape(-1, 4, 8, 8) y_train_policy = Variable(xp.array(batch_y_policy)).reshape(-1, 64) y_train_value = Variable(xp.array(batch_y_value)).reshape(-1, 1) return x_train, y_train_policy, y_train_value
最後にcreate-model-batchで動かす場合にGPUモードにしています。
elif len(args) > 2 and args[1] == 'create-model-batch': ### ...(省略)... gpu_device = 0 ### ...(省略)...
AWS Batchのコンピューティング環境で使うAMIとインスタンスを変更
AWS Batchで今回のAMIをGPUインスタンスを使うようにしました。
以下、コードです。
コンピューティング環境の定義のimageIdに今回のAMIをinstanceTypesにGPUインスタンスを指定します。今回はp2.xlargeにしました。
cat << EOF > compute-environment.spec.json
{
"computeEnvironmentName": "test-compute-environment",
"type": "MANAGED",
"state": "ENABLED",
"computeResources": {
"type": "EC2",
"minvCpus": 0,
"maxvCpus": 4,
"desiredvCpus": 0,
"instanceTypes": [
"p2.xlarge"
],
"imageId": "${AMI_ID}",
"subnets": ["${SUBNET_ID}"],
"securityGroupIds": ["${DEFAULT_SECURITY_GROUP_ID}"],
"instanceRole": "${INSTANCE_ROLE_ARN}"
},
"serviceRole": "${ROLE_SERVICE_ARN}"
}
EOF
COMPUTE_ENV=$(aws batch create-compute-environment --cli-input-json file://compute-environment.spec.json)
動かした結果
結論は前回と同じで４時間かかり学習時間が短縮することはありませんでした。
GPUが有効に動いていないことが原因ではないかと今回のAMI/インスタンス/docker imageでMNISTのサンプルを動かしてみましたがきちんと動作しているようです。GPUは有効だと思われます。
chainer/examples/mnist at v5 · chainer/chainer · GitHub
バグがある可能性や今回のケースではあまりGPUの効果がない可能性、バッチサイズを増やしたりパラメータを変更したら効果がある可能性、など疑っていますが調査中です... 🤔
まとめ
今回はAPV-MCTSのセルフプレイによる強化学習のAWS Batch環境をGPU化してみました。残念ながら学習時間は前回と同じで短縮できませんでしたがPackerの使い方を知ることができてGPU環境も作ることができました。引き続き調査と流石にそろそろコードが限界なのでリファクタリングをしたいです。
