且构网

分享程序员开发的那些事...
且构网 - 分享程序员编程开发的那些事

Tensorflow保存子类模型,该模型具有多个用于call()方法的参数

更新时间:2023-12-02 18:33:28

如果将输入打包在列表中,则可以成功导出模型.您还需要指定输入签名以导出模型,在这里对代码进行一些修改即可生效

You can export the model successfully if you package your inputs in a list. You also need to specify the input signatures to export your model, here your code with slight modifications which works

import tensorflow as tf
from tensorflow.keras.layers import Embedding, LSTM
import numpy as np
print('TensorFlow: ', tf.__version__)


vocab_size = 10000
LSTM_DIM = 256
BATCH_SIZE = 16
embedding_matrix = np.random.randn(vocab_size, 300)

class Encoder(tf.keras.Model):
    def __init__(self, vocab_size, embedding_matrix, n_units, batch_size):
        super(Encoder, self).__init__()
        self.n_units = n_units
        self.batch_size = batch_size

        self.embedding = Embedding(vocab_size, embedding_matrix.shape[1], weights=[embedding_matrix], trainable=True, mask_zero=True)
        self.lstm = LSTM(n_units, return_sequences=True, return_state=True, recurrent_initializer="glorot_uniform")


    @tf.function
    def call(self, inputs):
        input_utterence, initial_state = inputs
        input_embed = self.embedding(input_utterence)
        encoder_states, h1, c1 = self.lstm(input_embed, initial_state=initial_state)
        return encoder_states, h1, c1


    def create_initial_state(self):
        return tf.zeros((self.batch_size, self.n_units))

random_input = tf.random.uniform(shape=[BATCH_SIZE, 3], maxval=vocab_size, dtype=tf.int32)


encoder = Encoder(vocab_size, embedding_matrix, LSTM_DIM, BATCH_SIZE)

initial_state = [encoder.create_initial_state(), encoder.create_initial_state()]
_ = encoder([random_input, initial_state]) # required so that encoder.build is triggered


tf.saved_model.save(encoder, "encoder_model", signatures=encoder.call.get_concrete_function(
         [
             tf.TensorSpec(shape=[None, None], dtype=tf.int32, name='input_utterence'), 
             [
                 tf.TensorSpec(shape=[None, LSTM_DIM], dtype=tf.float32, name='initial_h'), 
                 tf.TensorSpec(shape=[None, LSTM_DIM], dtype=tf.float32, name='initial_c')
             ]
         ]))

loaded_model = tf.saved_model.load('encoder_model')
loaded_model([random_input, initial_state])

输出:

TensorFlow:  2.2.0-rc1
WARNING:tensorflow:From /home/dl_user/tf_stable/lib/python3.7/site-packages/tensorflow/python/ops/resource_variable_ops.py:1817: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.
Instructions for updating:
If using Keras pass *_constraint arguments to layers.
INFO:tensorflow:Assets written to: encoder_model/assets
(<tf.Tensor: shape=(16, 3, 256), dtype=float32, numpy=
 array([[[-0.06000457,  0.02422162, -0.05310762, ..., -0.01340707,
           0.12212028, -0.02747637],
         [ 0.13303193,  0.3119418 , -0.17995344, ..., -0.10185111,
           0.09568192,  0.06919193],
         [-0.08075664, -0.11490613, -0.20294832, ..., -0.14999194,
           0.02177649,  0.05538464]],

        [[-0.03792192, -0.08431012,  0.03687581, ..., -0.1768839 ,
          -0.10469476,  0.08730042],
         [-0.02956271,  0.43850696, -0.07400024, ...,  0.04097629,
           0.209705  ,  0.27194855],
         [ 0.02529916,  0.18367583, -0.11409087, ...,  0.0458075 ,
           0.2065246 ,  0.22976378]],

        [[ 0.04196627,  0.08302739,  0.02218204, ...,  0.07388053,
          -0.05696848, -0.31895265],
         [-0.00536443,  0.1566213 , -0.22412768, ...,  0.10560389,
           0.20187919, -0.1896591 ],
         [ 0.26364946,  0.13163888,  0.14586888, ...,  0.19517538,
           0.17677066, -0.40476215]],

        ...,

        [[ 0.10999472,  0.07398727,  0.23443945, ..., -0.1912791 ,
          -0.0195728 ,  0.11717851],
         [ 0.03978832,  0.07587367,  0.16567066, ..., -0.29463592,
           0.05950819,  0.0242265 ],
         [ 0.2505787 ,  0.15849623,  0.06635283, ..., -0.17969091,
           0.12549783, -0.11459641]],

        [[-0.20408148,  0.04629526,  0.00601436, ...,  0.21321473,
           0.04952445, -0.0129672 ],
         [-0.14671509,  0.2911171 ,  0.13047697, ..., -0.03531414,
          -0.16794083,  0.01575338],
         [-0.08337164,  0.08723269,  0.16235027, ...,  0.07919721,
           0.05701642,  0.15379705]],

        [[-0.2747393 ,  0.24351111, -0.05829309, ..., -0.00448833,
           0.07568972,  0.03978251],
         [-0.16282909, -0.04586324, -0.0054924 , ...,  0.11050001,
           0.1312355 ,  0.16555254],
         [ 0.07759799, -0.07308074, -0.10038756, ...,  0.18139914,
           0.07769153,  0.1375772 ]]], dtype=float32)>,
 <tf.Tensor: shape=(16, 256), dtype=float32, numpy=
 array([[-0.08075664, -0.11490613, -0.20294832, ..., -0.14999194,
          0.02177649,  0.05538464],
        [ 0.02529916,  0.18367583, -0.11409087, ...,  0.0458075 ,
          0.2065246 ,  0.22976378],
        [ 0.26364946,  0.13163888,  0.14586888, ...,  0.19517538,
          0.17677066, -0.40476215],
        ...,
        [ 0.2505787 ,  0.15849623,  0.06635283, ..., -0.17969091,
          0.12549783, -0.11459641],
        [-0.08337164,  0.08723269,  0.16235027, ...,  0.07919721,
          0.05701642,  0.15379705],
        [ 0.07759799, -0.07308074, -0.10038756, ...,  0.18139914,
          0.07769153,  0.1375772 ]], dtype=float32)>,
 <tf.Tensor: shape=(16, 256), dtype=float32, numpy=
 array([[-0.32829475, -0.18770668, -0.2956414 , ..., -0.2427501 ,
          0.03146099,  0.16033864],
        [ 0.05112522,  0.6664379 , -0.19836858, ...,  0.10015503,
          0.511694  ,  0.51550364],
        [ 0.3379809 ,  0.7145362 ,  0.22311993, ...,  0.372106  ,
          0.25914627, -0.81374717],
        ...,
        [ 0.36742535,  0.29009506,  0.13245934, ..., -0.4318537 ,
          0.26666188, -0.20086129],
        [-0.17384854,  0.22998339,  0.27335796, ...,  0.09973672,
          0.10726923,  0.47339764],
        [ 0.22148325, -0.11998752, -0.16339599, ...,  0.31903535,
          0.20365229,  0.28087002]], dtype=float32)>)