Search code examples
pythonmachine-learningkerasconv-neural-networkkeras-layer

keras setting trainable flag on pretrained model


Suppose I have a model

from tensorflow.keras.applications import DenseNet201

base_model = DenseNet201(input_tensor=Input(shape=basic_shape))

model = Sequential()
model.add(base_model)

model.add(Dense(400))
model.add(BatchNormalization())
model.add(ReLU())

model.add(Dense(50, activation='softmax'))

model.save('test.hdf5')

Then I load the saved model and try to make the last 40 layers of DenseNet201 trainable and the first 161 - non-trainable:

saved_model = load_model('test.hdf5')
cnt = 44
saved_model.trainable = False
  while cnt > 0:
      saved_model.layers[-cnt].trainable = True
      cnt -= 1

But this is not actually working because DenseNet201 is determined as a single layer and I just get index out of range error.

Layer (type)                 Output Shape              Param #   
=================================================================
densenet201 (Functional)     (None, 1000)              20242984  
_________________________________________________________________
dense (Dense)                (None, 400)               400400    
_________________________________________________________________
batch_normalization (BatchNo (None, 400)               1600      
_________________________________________________________________
re_lu (ReLU)                 (None, 400)               0         
_________________________________________________________________
dense_1 (Dense)              (None, 50)                20050     
=================================================================
Total params: 20,665,034
Trainable params: 4,490,090
Non-trainable params: 16,174,944

The question is how can I actually make the first 161 layers of DenseNet non-trainable and the last 40 layers trainable on a loaded model?


Solution

  • densenet201 (Functional) is a nested model, therefore you can access its layers the same way you access the layers of your 'topmost' model.

    saved_model.layers[0].layers
    

    where saved_model.layers[0] is a model with its own layers.

    In your loop, you need to access the layers like this

    saved_model.layers[0].layers[-cnt].trainable = True
    

    Update

    By default, the loaded model's layers are trainable (trainable=True), therefore you will need to set the bottom layers' trainable attribute to False instead.