混合精度でニューラルネットワークのトレーニングを高速化する事例：多層ニューラルネットワークの一つであるCNN(CNN: Convolution Neural Network)上のトレーニング時間の短縮
はじめに
この例では、多層ニューラルネットワークの一つであるCNN上で、画像を分類するためのトレーニング時間を短縮できることを示します。混合精度を使用することにより、FP16型で計算してもよいレイヤーとFP32型で計算してもよいレイヤーを自動的に識別し、画像の分類精度をほぼ落とすことなくトレーニング時間を短縮できます。今回GPUSOROBAN(A100, GPUメモリ：40GB)にて、混合精度の機能を利用してトレーニングを行った際、FP32型のみで計算したときに比べて約25%短縮することができました。参考記事１によると、Google Colab上のNVIDIA Tesla T4 GPUでは（同じモデルとバッチサイズを使用して）トレーニング時間を10エポックにわたって約33%短縮できる、という結果が得られています。
混合精度を利用するための必要条件：
ハードウェア側の条件：
・NVIDIA (Volta / Turing / Ampere) アーキテクチャ GPUのみ対応
・NVIDIA Driver release 418.xx+
・CUDA 10.1以上(*但し、NVIDIA Driver releaseに対応するもの)
・対応するcuDNNモジュールのインストール
・Tesla(Tesla V100, Tesla P4, Tesla P40 or Tesla P100)を利用される場合、NVIDIA driver release 384.111+ or 410を利用する必要があります。また、CUDA及びcuDNNについては、NVIDIA driverに対応するバーションを利用する必要があります。
ソフトウェア側の条件：
・TensorFlow version >=1.14.0-rc0
混合精度の仕組み
混合精度の仕組みはFP16型で計算してもよいレイヤーとFP32型で計算してもよいレイヤーを自動的に識別し、モデルのトレーニングを実施できることです。FP16型で計算したレイヤーにつき、アンダーフローが発生する場合、FP32型で計算するので、モデルの最適化精度を落とすことなく、トレーニング時間が短縮され、メモリ使用量が削減されます。
混合精度について、もっと知りたい場合、以下資料をお勧めします。
・Overview of Automatic Mixed Precision for Deep Learning
・NVIDIA Mixed Precision Training Documentation
・NVIDIA Deep Learning Performance Guide
・Information about NVIDIA Tensor Cores
・Post on TensorFlow blog explaining Automatic Mixed Precision
# Copyright 2019 NVIDIA Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
下記のセクションでは必要なライブラリーがインストールされているか、そして、Tensor Coreが入っているのかの確認です。
import time
import numpy as np
import tensorflow.compat.v2 as tf
tf.enable_v2_behavior()
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import *
from tensorflow.python.client import device_lib
def check_tensor_core_gpu_present():
local_device_protos = device_lib.list_local_devices()
for line in local_device_protos:
if "compute capability" in str(line):
compute_capability = float(line.physical_device_desc.split("compute capability: ")[-1])
if compute_capability>=7.0:
return True
print("TensorFlow version is", tf.__version__)
try:
# check and assert TensorFlow >= 1.14
tf_version_list = tf.__version__.split(".")
if int(tf_version_list[0]) < 2:
assert int(tf_version_list[1]) >= 14
except:
print("TensorFlow 1.14.0 or newer is required.")
print("Tensor Core GPU Present:", check_tensor_core_gpu_present())
if check_tensor_core_gpu_present():
pass
else:
!nvidia-smi
assert check_tensor_core_gpu_present() == True
データセットのインポート
CIFAR10 の画像データセットを tf.keras.datasets からインポートします。
# The data, split between train and test sets
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
num_classes = np.max(y_train) + 1
# Convert class vectors to binary class matrices
y_train = tf.keras.utils.to_categorical(y_train, num_classes)
y_test = tf.keras.utils.to_categorical(y_test, num_classes)
画像の前処理として、画像の値をレンジ0 ～ 255からレンジ0 ～ 1に正規化します。
def normalize(ndarray):
ndarray = ndarray.astype("float32")
ndarray = ndarray/255.0
return ndarray
x_train = normalize(x_train)
x_test = normalize(x_test)
モデルの定義
画像の分類ができるように単純なCNNを返すための再利用可能なヘルパー関数を定義します。
def create_model(num_classes=10):
"""
Returns a simple CNN suitable for classifiying images from CIFAR10
"""
# model parameters
act = "relu"
pad = "same"
ini = "he_uniform"
model = tf.keras.models.Sequential([
Conv2D(128, (3, 3), activation=act, padding=pad, kernel_initializer=ini,
input_shape=(32,32,3)),
Conv2D(256, (3, 3), activation=act, padding=pad, kernel_initializer=ini),
Conv2D(256, (3, 3), activation=act, padding=pad, kernel_initializer=ini),
Conv2D(256, (3, 3), activation=act, padding=pad, kernel_initializer=ini),
MaxPooling2D(pool_size=(2,2)),
Conv2D(256, (3, 3), activation=act, padding=pad, kernel_initializer=ini),
Conv2D(256, (3, 3), activation=act, padding=pad, kernel_initializer=ini),
Conv2D(512, (3, 3), activation=act, padding=pad, kernel_initializer=ini),
Conv2D(512, (3, 3), activation=act, padding=pad, kernel_initializer=ini),
MaxPooling2D(pool_size=(2,2)),
Conv2D(256, (3, 3), activation=act, padding=pad, kernel_initializer=ini),
Conv2D(256, (3, 3), activation=act, padding=pad, kernel_initializer=ini),
Conv2D(256, (3, 3), activation=act, padding=pad, kernel_initializer=ini),
Conv2D(128, (3, 3), activation=act, padding=pad, kernel_initializer=ini),
MaxPooling2D(pool_size=(4,4)),
Flatten(),
BatchNormalization(),
Dense(512, activation='relu'),
Dense(num_classes, activation="softmax")
])
return model
model = create_model(num_classes)
model.summary()
モデルのトレーニング
混合精度を用いない場合と用いる場合の比較を下記の関数(train_model)で、実施します。参考記事１では "optimizer = tensorflow.compat.v1.train.experimental.enable_mixed_precision_graph_rewrite(optimizer)" が使われております。TensorFlow2.4.0以降ではこのAPI自体がもはや実験的ではなくなるため、下記のコードに変更され、より高速にトレーニングすることができます。
optimizer = tf.compat.v1.mixed_precision.enable_mixed_precision_graph_rewrite(optimizer) *このような変更するだけ、時間短縮の度合いは34%から48％まで高まります。
# training parameters
BATCH_SIZE = 320
N_EPOCHS = 10
opt = tf.keras.optimizers.SGD(learning_rate=0.02, momentum=0.5)
def train_model(mixed_precision, optimizer):
"""
Trains a CNN to classify images on CIFAR10,
and returns the training and classification performance
Args:
mixed_precision: `True` or `False`
optimizer: An instance of `tf.keras.optimizers.Optimizer`
"""
model = create_model(num_classes)
if mixed_precision:
import tensorflow
optimizer = tensorflow.compat.v1.train.experimental.enable_mixed_precision_graph_rewrite(optimizer)
model.compile(loss="categorical_crossentropy",
optimizer=optimizer,
metrics=["accuracy"])
train_start = time.time()
train_log = model.fit(x_train, y_train,
batch_size=BATCH_SIZE,
epochs=N_EPOCHS,
use_multiprocessing=True,
workers=2)
## 参考記事２をご参照ください。
score = model.evaluate(x_test, y_test, verbose=0)
##
train_end = time.time()
results = {"test_loss": score[0],
"test_acc": score[1],
"train_time": train_end-train_start,
"train_log": train_log}
return results
*参考記事１のコードではscore = model.evaluate(x_test, y_test, verbose=0)という定義がなされていないバグがあり、参考記事２の通りで修正することで解消されます。
FP32型の場合の計算を下記の通り、行います。
fp32_results = train_model(mixed_precision=False, optimizer=opt)
test_acc = round(fp32_results["test_acc"]*100, 1)
train_time = round(fp32_results["train_time"], 1)
print(test_acc, "% achieved in", train_time, "seconds")
下記のコードでは10秒間ほど、GPUを休憩させ、クールダウンさせます。
# to ensure accuracy of timing benchmark
# we give the GPU 10 seconds to cool down
tf.keras.backend.clear_session()
time.sleep(10)
混合精度の場合の計算を下記の通り、行います。
mp_results = train_model(mixed_precision=True, optimizer=opt)
test_acc = round(mp_results["test_acc"]*100, 1)
train_time = round(mp_results["train_time"], 1)
print(test_acc, "% achieved in", train_time, "seconds")
import matplotlib.pyplot as plt
%matplotlib inline
plt.plot(fp32_results["train_log"].history["loss"], label="FP32")
plt.plot(mp_results["train_log"].history["loss"], label="Mixed Precision")
plt.title("Performance Comparison")
plt.ylabel("Training Loss")
plt.xlabel("Epoch")
plt.legend()
plt.show()
