我想建立自己的基础模型并使用大数据集对其进行训练。训练后,我保存了基础模型。我有另一个自定义模型,我想从基础模型中加载前两层的权重。我应该如何在 Tensorflow 2.1.0 中实现它,谢谢。
示例代码:
import os
os.environ["CUDA_VISIBLE_DEVICES"]=""
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
class BaseModel():
def __init__(self):
inputs = keras.Input(shape=(32, 32, 3))
x = inputs
x = layers.Conv2D(32, 3, padding='same', activation=tf.nn.relu)(x)
x = layers.MaxPool2D()(x)
x = layers.Conv2D(64, 3, padding='same', activation=tf.nn.relu)(x)
x = layers.Flatten()(x)
x = layers.Dense(500, activation=tf.nn.relu)(x)
outputs = layers.Dense(1000, activation=tf.nn.softmax)(x)
self.model = keras.Model(inputs=inputs, outputs=outputs)
def __call__(self, inputs):
return self.model(inputs)
bm = BaseModel() # the model for pretraining
bm.model.save_weights('base_model') # save the pretrained model
class MyModel():
def __init__(self):
inputs = keras.Input(shape=(32, 32, 3))
x = inputs
x = layers.Conv2D(32, 3, padding='same', activation=tf.nn.relu)(x)
x = layers.MaxPool2D()(x)
x = layers.Conv2D(64, 3, padding='same', activation=tf.nn.relu)(x)
x = layers.Conv2D(128, 3, padding='same', activation=tf.nn.relu)(x)
x = layers.Flatten()(x)
x = layers.Dense(1000, activation=tf.nn.relu)(x)
outputs = layers.Dense(10, activation=tf.nn.softmax)(x)
self.model = keras.Model(inputs=inputs, outputs=outputs)
def __call__(self, inputs):
return self.model(inputs)
mm = MyModel() # the model for my customized applications
mm.model.load_weights('base_model') # load the pretrained model with the first two conv layers
# further fine-tuning or transfer learning