TF 2.2 Saved_model из keras с настраиваемыми подписями и предварительной обработкой - PullRequest
Новая
       91

TF 2.2 Saved_model из keras с настраиваемыми подписями и предварительной обработкой

1 голос
/ 19 июня 2020

Я пытаюсь использовать tf.saved_model.save после обучения модели трансформатора, чтобы развернуть его. Моя модель имеет несколько входов и выходов. Если я использую функцию saved_model для обслуживания, я добавляю некоторую проблему с формой ввода с первым вводом, а второй ввод не отображается, когда я использую функцию saved_model_cli show. Я нашел способ решить эту проблему, обернув основной блок трансформатора модулем модели, а затем сохранив модель. 

def build_model(MAX_LEN, transformer_layer):
    inp = Input(shape=(MAX_LEN,), dtype=tf.int32, name="input_word_ids")
    inp2 = Input(shape=(MAX_LEN,), dtype=tf.int32, name="attention_mask")
    sequence_output, pooled_output = transformer_layer(inp, attention_mask = inp2)
    out = Dense(1, activation='sigmoid', name='outputs')(pooled_output)
    model = Model(inputs=[inp], outputs=[out])
    return model

Но, как я вижу, в документации есть другой способ, который заключается в использовании подписей для указать входы / выходы во время tf.saved_model.save. Для этого нам нужно использовать класс tf.Module, но я не понял, как его использовать в случае, если у нас есть несколько входов / выходов точно (как понять модуль, что тензор входов должен быть связан с этими входами для модель?) . Кто-нибудь знает, как это сделать вторым методом? Более того, можем ли мы сделать предварительную обработку данных через подпись? https://www.tensorflow.org/guide/saved_model

class CustomModule(tf.Module):

  def __init__(self):
    super(CustomModule, self).__init__()
    self.v = tf.Variable(1.)

  @tf.function
  def __call__(self, x):
    print('Tracing with', x)
    return x * self.v

  @tf.function(input_signature=[tf.TensorSpec([], tf.float32)])
  def mutate(self, new_v):
    self.v.assign(new_v)

module = CustomModule()
...