Google Colaboratoryで動いたJupyter上のTensorflow/KerasコードがAWS SageMakerで動かない

2023/04/04

当サイトではアフィリエイト広告を利用しています

Google Colaboratoryで動いていたJupyter NotebookのコードをAWS Sagemakerのノートブックインスタンスで動かしてみたところ、Shape must be at least rank 3 but is rank 2 for '{{node BiasAdd}} = BiasAdd[T=DT_FLOAT, data_format="NCHW"](add, bias)' with input shapes: [?,128], [128].のエラーが出て動作しませんでした。

原因を調べて対処方法をまとめましたので記録しておきます。

現在、改めてDeepLearningを自分で使って実装する必要が出てきたので、Pythonによるディープラーニングを使ってpython + tensorflow + kerasの組み合わせでの実装を学習しています。

書籍中のCh11のサンプルコードにおいて、Google Colaboratoryで動いたJupyterコードがAWS Sagemakerでは動作せず、次のようなエラーを出力しました。

Shape must be at least rank 3 but is rank 2 for '{{node BiasAdd}} = BiasAdd[T=DT_FLOAT, data_format="NCHW"](add, bias)' with input shapes: [?,128], [128].

結論を書くと、Kerasのbackendのデフォルト値がchannels_first/channels_lastでColaboratoryとSagemakerで違っていることが原因だとわかりました。これの差でTensorのRank拡張時の挙動に差が出てエラーとなったようです。

環境毎の違い

バージョン確認、チャネル確認用コードをそれぞれの環境で走らせてみました。

import tensorflow as tf
print(tf.__version__)

from tensorflow import keras
print(keras.__version__)
print(keras.backend.image_data_format())

tensorflow,kerasのバージョンを出力し、同時にkerasのbackendについてchannels_first/channels_lastの違いを出力しています。

Google Colaboratory

Google Colaboratoryの実行結果は以下でした。

2.11.0
2.11.0
channels_last

AWS Sagemaker

kernelはconda_tensorflow2_p310を利用しています。この時の出力は以下でした。

2.11.0
2.11.0
channels_first

比較

tensorflowのバージョンは同じでも、channelの値が違うことがわかります。

channels_first: N,C,H,W
channels_last: N,H,W,C

デフォルト値はchannels_lastですが、AWSのカーネルではchannels_firstへ切り替えているようです。

N: the Number of images in batch
C: Channel
H: Height
W: Width

どうもCuDNNでNvidiaGPUを使う場合はchannels_firstの方が少し速いようで、Sagemakerも設定を変えていたのかもしれません。

対処方法

以下のようにすればBackendを変更できます。(公式Doc)

tf.keras.backend.set_image_data_format("channels_last")

AWS Sagemakerでの対処

SagemakerでJupyterの最初に以下のようなコードを挿入し、Colaboratoryと同じchannels_lastに変更することで、書籍のサンプルコードが動作するようになりました。

import tensorflow as tf
print(f"Tensorflow Version: {tf.__version__}")

from tensorflow import keras
print(f"Keras Version: {keras.__version__}")
print(f"Keras Backend data_format: {keras.backend.image_data_format()}")

# Google Colaboratoryに揃えるためにdata_formatを"channels_last"に変更しておく
tf.keras.backend.set_image_data_format("channels_last")
print(f"Keras Backend data_format: {keras.backend.image_data_format()}")

実行時の出力は以下です。

Tensorflow Version: 2.11.0
Keras Version: 2.11.0
Keras Backend data_format: channels_first
Keras Backend data_format: channels_last

その他

色々状況を調べたのち、自身でAWS Sagemaker上にColaboratoryと同じバージョンでConda環境構築することも考えましたが、tensorflowを一からインストールするとEBSの容量をかなり必要とします。

EBSのデフォルト値は最低限の5GBでしたが、cuda関連やtensorflow2をインストールしてみたところ、5GB超えて途中でインストール止まってしまいました。

最終的には上記のdata_formatで解決できたので、EBS容量を大きくしたりライフサイクル設定など面倒なところに踏み込まずに済んだようで助かりました。

まとめ

Google ColaboratoryとAWS Sagemakerでデフォルト値が違う
これが原因でRank拡張時の動作が異なってエラーが出ていた
tf.keras.backend.set_image_data_format()で対処可能