TensorflowのKerasがVS Codeで補完されるようにする

VS CodeでTensorflowのKerasをimportすると、Lasy Loadのせいで補完が効かない。そのため、補完を効かせるための方法をメモ。

環境

  • VS Code: 1.85.0
  • Python: 3.10.3
  • Tensorflow: 2.16.1
  • Keras: 3.1.1

問題点

Tensorflowで普通にKerasをimportする場合、以下のように書く。

import tensorflow.keras as keras

しかし、tensorflow.kerasは実行時にLazy Loadされるため、このままだとPylanceに"不明なインポートシンボル"で怒られる。

Tensorflow側の実装のこのあたりを見ると、Keras v3であればkeras._tf_keras.kerasが実体になる様子。

解決策

以下のように書くことで、コーディング時はkeras._tf_keras.kerasを直接参照し、実行時はtensorflow.kerasが使われるようになる。 (補完のためのハックがコードに染み出してしまうのが嫌なところ…)

import tensorflow as tf
if not hasattr(tf, "keras"):
    import keras._tf_keras as keras
else:
    import tensorflow.keras as keras