TensorFlow, Kerasでレイヤー、モデルのtrainable属性を設定（Freeze / Unfreeze）
TensorFlow, Kerasのレイヤーやモデルのtrainable属性で、そのレイヤーまたはモデルを訓練（学習）するかしないか、すなわち、訓練時にパラメータ（カーネルの重みやバイアスなど）を更新するかどうかを設定できる。
レイヤーやモデルを訓練対象から除外することを「freeze（凍結）」、freezeしたレイヤーやモデルを再び訓練対象とすることを「Unfreeze（解凍、凍結解除）」と呼ぶ。
ここでは、以下の内容について説明する。
trainable属性の変更後はcompile()が必要- レイヤーの
trainable属性を変更 - モデルの
trainable属性を変更 - ネストしたモデルの場合
なお、BatchNormalization層におけるtrainable属性の振る舞いは特殊なので注意。以下の記事を参照。
以下のサンプルコードのTensorFlowのバージョンは2.1.0。TensorFlowに統合されたKerasを使う。
import tensorflow as tf
print(tf.__version__)
# 2.1.0
source: tf_keras_trainable_basics.py
trainable属性の変更後はcompile()が必要
モデルの訓練プロセス（オプティマイザー、損失関数、評価関数）の設定にはcompile()を使う。
trainable属性の設定はcompile()によって有効化される。compile()してからtrainable属性を変更した場合は、再度compile()する必要がある。
インスタンス化後にレイヤーの
trainableプロパティにTrueかFalseを設定することができます．設定の有効化のためには，trainableプロパティの変更後のモデルでcompile()を呼ぶ必要があります． FAQ - レイヤーを"freeze"するには？ - Keras Documentation
trainable属性を変更して繰り返し訓練を行うとき、例えば学習済みモデルのファインチューニングのときなどは要注意。
レイヤーのtrainable属性を変更
以下のモデルを例とする。ただの例なので、なにか意味のあるモデルではない。
model = tf.keras.Sequential([
tf.keras.layers.Dense(100, name='Layer_0', input_shape=(1000,)),
tf.keras.layers.Dense(10, name='Layer_1'),
tf.keras.layers.Dense(1, name='Layer_2')
], name='Sequential')
model.summary()
# Model: "Sequential"
# _________________________________________________________________
# Layer (type) Output Shape Param #
# =================================================================
# Layer_0 (Dense) (None, 100) 100100
# _________________________________________________________________
# Layer_1 (Dense) (None, 10) 1010
# _________________________________________________________________
# Layer_2 (Dense) (None, 1) 11
# =================================================================
# Total params: 101,121
# Trainable params: 101,121
# Non-trainable params: 0
# _________________________________________________________________
source: tf_keras_trainable_basics.py
以降、summary()の出力でTrainable params（訓練対象のパラメータ数）、Non-trainable params（訓練対象でないパラメータ数）を確認する。それぞれを値として取得する方法もある。
対象のレイヤーのオブジェクトを指定し、trainable属性をTrueまたはFalseに設定する。
model.layers[1].trainable = False
model.get_layer('Layer_2').trainable = False
for l in model.layers:
print(l.name, l.trainable)
# Layer_0 True
# Layer_1 False
# Layer_2 False
model.summary()
# Model: "Sequential"
# _________________________________________________________________
# Layer (type) Output Shape Param #
# =================================================================
# Layer_0 (Dense) (None, 100) 100100
# _________________________________________________________________
# Layer_1 (Dense) (None, 10) 1010
# _________________________________________________________________
# Layer_2 (Dense) (None, 1) 11
# =================================================================
# Total params: 101,121
# Trainable params: 100,100
# Non-trainable params: 1,021
# _________________________________________________________________
source: tf_keras_trainable_basics.py
for文を利用して複数のレイヤーを一括変更することも可能。
for l in model.layers[1:]:
l.trainable = True
for l in model.layers:
print(l.name, l.trainable)
# Layer_0 True
# Layer_1 True
# Layer_2 True
model.summary()
# Model: "Sequential"
# _________________________________________________________________
# Layer (type) Output Shape Param #
# =================================================================
# Layer_0 (Dense) (None, 100) 100100
# _________________________________________________________________
# Layer_1 (Dense) (None, 10) 1010
# _________________________________________________________________
# Layer_2 (Dense) (None, 1) 11
# =================================================================
# Total params: 101,121
# Trainable params: 101,121
# Non-trainable params: 0
# _________________________________________________________________
source: tf_keras_trainable_basics.py
レイヤーオブジェクトの取得については以下の記事を参照。インデックス・名前で指定するほか、条件を指定して取得することもできる。
モデル生成時にレイヤーの引数trainableで初期値を設定することも可能。
model_1 = tf.keras.Sequential([
tf.keras.layers.Dense(100, name='Layer_0', input_shape=(1000,)),
tf.keras.layers.Dense(10, name='Layer_1', trainable=False),
tf.keras.layers.Dense(1, name='Layer_2')
], name='Sequential_1')
model_1.summary()
# Model: "Sequential_1"
# _________________________________________________________________
# Layer (type) Output Shape Param #
# =================================================================
# Layer_0 (Dense) (None, 100) 100100
# _________________________________________________________________
# Layer_1 (Dense) (None, 10) 1010
# _________________________________________________________________
# Layer_2 (Dense) (None, 1) 11
# =================================================================
# Total params: 101,121
# Trainable params: 100,111
# Non-trainable params: 1,010
# _________________________________________________________________
source: tf_keras_trainable_basics.py
モデルのtrainable属性を変更
Kerasでは、モデルtf.keras.Modelはレイヤーの基底クラスであるtf.keras.layers.Layerのサブクラスであり、モデルもtrainable属性を持つ。
print(issubclass(tf.keras.Model, tf.keras.layers.Layer))
# True
source: tf_keras_trainable_basics.py
モデルのtrainable属性を変更すると、そのモデルに含まれるレイヤーのtrainable属性も再帰的に変更される。
model.trainable = False
for l in model.layers:
print(l.name, l.trainable)
# Layer_0 False
# Layer_1 False
# Layer_2 False
model.summary()
# Model: "Sequential"
# _________________________________________________________________
# Layer (type) Output Shape Param #
# =================================================================
# Layer_0 (Dense) (None, 100) 100100
# _________________________________________________________________
# Layer_1 (Dense) (None, 10) 1010
# _________________________________________________________________
# Layer_2 (Dense) (None, 1) 11
# =================================================================
# Total params: 101,121
# Trainable params: 0
# Non-trainable params: 101,121
# _________________________________________________________________
source: tf_keras_trainable_basics.py
モデルのtrainable属性がFalseの場合、内部のレイヤーのtrainable属性によらずモデル全体がFreezeされる。レイヤーのtrainable属性をTrueとしても訓練対象にならないので注意。
Trainable paramsが0であることに注目。
model.layers[1].trainable = True
for l in model.layers:
print(l.name, l.trainable)
# Layer_0 False
# Layer_1 True
# Layer_2 False
model.summary()
# Model: "Sequential"
# _________________________________________________________________
# Layer (type) Output Shape Param #
# =================================================================
# Layer_0 (Dense) (None, 100) 100100
# _________________________________________________________________
# Layer_1 (Dense) (None, 10) 1010
# _________________________________________________________________
# Layer_2 (Dense) (None, 1) 11
# =================================================================
# Total params: 101,121
# Trainable params: 0
# Non-trainable params: 101,121
# _________________________________________________________________
source: tf_keras_trainable_basics.py
なお、上の状態からモデルのtrainable属性をTrueとすると、上述のように内部のレイヤーのtrainable属性も再帰的にTrueとなる。
model.trainable = True
for l in model.layers:
print(l.name, l.trainable)
# Layer_0 True
# Layer_1 True
# Layer_2 True
source: tf_keras_trainable_basics.py
特定のレイヤーのみを訓練対象としたい場合は、全てのレイヤーのtrainable属性をFalseとしてから所望のレイヤーのtrainable属性をTrueとする。
for l in model.layers:
l.trainable = False
model.layers[1].trainable = True
for l in model.layers:
print(l.name, l.trainable)
# Layer_0 False
# Layer_1 True
# Layer_2 False
model.summary()
# Model: "Sequential"
# _________________________________________________________________
# Layer (type) Output Shape Param #
# =================================================================
# Layer_0 (Dense) (None, 100) 100100
# _________________________________________________________________
# Layer_1 (Dense) (None, 10) 1010
# _________________________________________________________________
# Layer_2 (Dense) (None, 1) 11
# =================================================================
# Total params: 101,121
# Trainable params: 1,010
# Non-trainable params: 100,111
# _________________________________________________________________
source: tf_keras_trainable_basics.py
ネストしたモデルの場合
上述のように、モデルtf.keras.Modelはレイヤーの基底クラスであるtf.keras.layers.Layerのサブクラスであるため、モデルを一つのレイヤーとして扱うことができる。
ネストした（入れ子になった）モデルを例とする。
inner_model = tf.keras.Sequential([
tf.keras.layers.Dense(100, name='Layer_in_0', input_shape=(1000,)),
tf.keras.layers.Dense(10, name='Layer_in_1')
], name='Inner_model')
inner_model.summary()
# Model: "Inner_model"
# _________________________________________________________________
# Layer (type) Output Shape Param #
# =================================================================
# Layer_in_0 (Dense) (None, 100) 100100
# _________________________________________________________________
# Layer_in_1 (Dense) (None, 10) 1010
# =================================================================
# Total params: 101,110
# Trainable params: 101,110
# Non-trainable params: 0
# _________________________________________________________________
outer_model = tf.keras.Sequential([
inner_model,
tf.keras.layers.Dense(1, name='Layer_out_1')
], name='Outer_model')
outer_model.summary()
# Model: "Outer_model"
# _________________________________________________________________
# Layer (type) Output Shape Param #
# =================================================================
# Inner_model (Sequential) (None, 10) 101110
# _________________________________________________________________
# Layer_out_1 (Dense) (None, 1) 11
# =================================================================
# Total params: 101,121
# Trainable params: 101,121
# Non-trainable params: 0
# _________________________________________________________________
print(outer_model.layers[0] is inner_model)
# True
内側のモデル内のレイヤーのtrainable属性も個別に変更できる
inner_model.layers[1].trainable = False
for l in inner_model.layers:
print(l.name, l.trainable)
# Layer_in_0 True
# Layer_in_1 False
for l in outer_model.layers:
print(l.name, l.trainable)
# Inner_model True
# Layer_out_1 True
outer_model.summary()
# Model: "Outer_model"
# _________________________________________________________________
# Layer (type) Output Shape Param #
# =================================================================
# Inner_model (Sequential) (None, 10) 101110
# _________________________________________________________________
# Layer_out_1 (Dense) (None, 1) 11
# =================================================================
# Total params: 101,121
# Trainable params: 100,111
# Non-trainable params: 1,010
# _________________________________________________________________
外側のモデルのtrainable属性を変更すると、内側のモデル内のレイヤーのtrainable属性まですべて再帰的に変更される。
outer_model.trainable = True
for l in inner_model.layers:
print(l.name, l.trainable)
# Layer_in_0 True
# Layer_in_1 True
for l in outer_model.layers:
print(l.name, l.trainable)
# Inner_model True
# Layer_out_1 True
outer_model.summary()
# Model: "Outer_model"
# _________________________________________________________________
# Layer (type) Output Shape Param #
# =================================================================
# Inner_model (Sequential) (None, 10) 101110
# _________________________________________________________________
# Layer_out_1 (Dense) (None, 1) 11
# =================================================================
# Total params: 101,121
# Trainable params: 101,121
# Non-trainable params: 0
# _________________________________________________________________
上述のように、モデルのtrainable属性がFalseの場合、内部のレイヤーのtrainable属性によらずモデル全体がFreezeされる。
内側のモデルinner_modelのtrainable属性をFalseとすると、その中のレイヤーのtrainable属性をTrueとしても訓練対象にならないので注意。
inner_model.trainable = False
inner_model.layers[1].trainable = True
for l in inner_model.layers:
print(l.name, l.trainable)
# Layer_in_0 False
# Layer_in_1 True
for l in outer_model.layers:
print(l.name, l.trainable)
# Inner_model False
# Layer_out_1 True
outer_model.summary()
# Model: "Outer_model"
# _________________________________________________________________
# Layer (type) Output Shape Param #
# =================================================================
# Inner_model (Sequential) (None, 10) 101110
# _________________________________________________________________
# Layer_out_1 (Dense) (None, 1) 11
# =================================================================
# Total params: 101,121
# Trainable params: 11
# Non-trainable params: 101,110
# _________________________________________________________________
Functional APIで既存のモデルを利用
以下のように、Functional APIで既存のモデルを使って新たなモデルを生成できる。
outer_model.trainable = True
functional_model = tf.keras.Model(
inputs=inner_model.input,
outputs=tf.keras.layers.Dense(1, name='Layer_2')(inner_model.output),
name='Functional_model'
)
functional_model.summary()
# Model: "Functional_model"
# _________________________________________________________________
# Layer (type) Output Shape Param #
# =================================================================
# Layer_in_0_input (InputLayer [(None, 1000)] 0
# _________________________________________________________________
# Layer_in_0 (Dense) (None, 100) 100100
# _________________________________________________________________
# Layer_in_1 (Dense) (None, 10) 1010
# _________________________________________________________________
# Layer_2 (Dense) (None, 1) 11
# =================================================================
# Total params: 101,121
# Trainable params: 101,121
# Non-trainable params: 0
# _________________________________________________________________
上のSequential APIの例とは異なり、モデルが一つのレイヤーとして扱われるわけではないが、元のモデルのレイヤーと新たなモデルのレイヤーは同じオブジェクトを指している。
print(inner_model.layers[0] is functional_model.layers[1])
# True
したがって、元のモデル（inner_model）のtrainable属性を変更すると、その中のレイヤーのtrainable属性も再帰的に変更され、新たなモデル（functional_model）にも影響する。
inner_model.trainable = False
for l in functional_model.layers:
print(l.name, l.trainable)
# Layer_in_0_input False
# Layer_in_0 False
# Layer_in_1 False
# Layer_2 True
functional_model.summary()
# Model: "Functional_model"
# _________________________________________________________________
# Layer (type) Output Shape Param #
# =================================================================
# Layer_in_0_input (InputLayer [(None, 1000)] 0
# _________________________________________________________________
# Layer_in_0 (Dense) (None, 100) 100100
# _________________________________________________________________
# Layer_in_1 (Dense) (None, 10) 1010
# _________________________________________________________________
# Layer_2 (Dense) (None, 1) 11
# =================================================================
# Total params: 101,121
# Trainable params: 11
# Non-trainable params: 101,110
# _________________________________________________________________
ただし、新たなモデル（functional_model）において元のモデル（inner_model）を一つのレイヤーとして扱っているわけではないので、元のモデルのtrainable属性がFalseのままでも、内部のレイヤーのtrainable属性を個別にTrueにすると、新たなモデルの中では訓練対象となる。
functional_model.layers[1].trainable = True
for l in functional_model.layers:
print(l.name, l.trainable)
# Layer_in_0_input False
# Layer_in_0 True
# Layer_in_1 False
# Layer_2 True
functional_model.summary()
# Model: "Functional_model"
# _________________________________________________________________
# Layer (type) Output Shape Param #
# =================================================================
# Layer_in_0_input (InputLayer [(None, 1000)] 0
# _________________________________________________________________
# Layer_in_0 (Dense) (None, 100) 100100
# _________________________________________________________________
# Layer_in_1 (Dense) (None, 10) 1010
# _________________________________________________________________
# Layer_2 (Dense) (None, 1) 11
# =================================================================
# Total params: 101,121
# Trainable params: 100,111
# Non-trainable params: 1,010
# _________________________________________________________________