Tensorflow 2.8 搭建 Seq2Seq 模型完成翻譯任務

語言: CN / TW / HK

本文正在參加「金石計劃 . 瓜分6萬現金大獎」

前言

本文使用 cpu 版本的 TensorFlow 2.8 版本完成西班牙文到英文的翻譯任務,我們假定讀者已經熟悉了 Seq2Seq 的模型,如果還不瞭解可以移步看我之前寫的文章,或者看相關論文:

  • 《Effective Approaches to Attention-based Neural Machine Translation》 (Luong et al., 2015)
  • https://juejin.cn/post/6973930281728213006

建議安裝好相關的工具包:

tensorflow-text==2.10
einops

大綱

  1. 獲取數據
  2. 處理數據
  3. 搭建 Encoder
  4. 搭建 Attention
  5. 搭建 Decoder
  6. 搭建完整的 Translator 模型
  7. 編譯、訓練
  8. 推理
  9. 保存和加載模型

實現

1. 獲取數據

(1)本文使用了一份西班牙文轉英文的數據,每一行都是一個樣本,每個樣本有一個西班牙文和對應的英文翻譯,兩者中間由一個水平製表符進行分割。

(2)我們可以使用 TensorFlow 的內置函數來從網絡上下載本文所用到的數據到本地,一般會下載到本地的 C:\Users\《用户名》.keras\datasets\spa-eng.zip 路徑下面。

(3)我們首使用 utf-8 的編碼格式讀取磁盤中的文件到內存,然後將每一行的樣本用水平製表符切割,將西班牙文作為我們的輸入,將英文作為我們的輸出目標。

(4)隨機選取 80% 的數據作為我們的訓練集,剩下的 20% 的數據當做驗證集。

(5)我們隨機選取了一個樣本的輸入和目標進行顯示,可以看到在模型的 Encoder 部分使用的是西班牙文,在模型的 Decoder 部分使用的是英文。

import pathlib
import numpy as np
import typing
from typing import Any, Tuple
import einops
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import tensorflow as tf
import tensorflow_text as tf_text
path_to_zip = tf.keras.utils.get_file('spa-eng.zip', origin='http://storage.googleapis.com/download.tensorflow.org/data/spa-eng.zip', extract=True)
path_to_file = pathlib.Path(path_to_zip).parent/'spa-eng/spa.txt'
BUFFER_SIZE = len(context_raw)
BATCH_SIZE = 128
def load_data(path):
    text = path.read_text(encoding='utf-8')
    lines = text.splitlines()
    pairs = [line.split('\t') for line in lines]
    context = np.array([context for target, context in pairs])
    target = np.array([target for target, context in pairs])
    return target, context
target_raw, context_raw = load_data(path_to_file)
is_train = np.random.uniform(size=(len(target_raw),)) < 0.8
train_raw = ( tf.data.Dataset.from_tensor_slices((context_raw[is_train], target_raw[is_train])).shuffle(BUFFER_SIZE).batch(BATCH_SIZE))
val_raw = ( tf.data.Dataset.from_tensor_slices((context_raw[~is_train], target_raw[~is_train])).shuffle(BUFFER_SIZE).batch(BATCH_SIZE))
for example_context, example_target in train_raw.take(1):
    for a,b in zip(example_context[:1], example_target[:1]):
        print(a)
        print(b)
    break

樣本結果打印:

tf.Tensor(b'A los gatos no les gusta estar mojados.', shape=(), dtype=string)
tf.Tensor(b'Cats dislike being wet.', shape=(), dtype=string)

2. 處理數據

(1)因為使用到不同的語言不同,可能涉及到不同的編碼問題,所以我們要使用 TensorFlow 內置的函數,將輸入和目標的所有文本都進行編碼的標準化,統一使用 utf-8 ,並且將文本中除了字母、空格、以及若干個指定的標點符號之外的字符都進行剔除,並且在輸入和目標文本的開始和末端加入 [START] 和 [END] 來表示句子的開始和結束。

(2)因為輸入和目標需要維護不同的詞典,所以我們對 token 進行整數映射的時候要維護一個輸入詞典和一個目標詞典,並且兩個詞典都要有 token -> int 和 int -> token 的映射。

def standardize(text):
    text = tf_text.normalize_utf8(text, 'NFKD')
    text = tf.strings.lower(text)
    text = tf.strings.regex_replace(text, '[^ a-z.?!,¿]', '')
    text = tf.strings.regex_replace(text, '[.?!,¿]', r' \0 ')
    text = tf.strings.strip(text)
    text = tf.strings.join(['[START]', text, '[END]'], separator=' ')
    return text
max_vocab_size = 6000
context_processor = tf.keras.layers.TextVectorization(standardize=standardize, max_tokens=max_vocab_size, ragged=True)
context_processor.adapt(train_raw.map(lambda context, target: context))
target_processor = tf.keras.layers.TextVectorization( standardize=standardize,  max_tokens=max_vocab_size, ragged=True)
target_processor.adapt(train_raw.map(lambda context, target: target))
print(context_processor.get_vocabulary()[:10])
print(target_processor.get_vocabulary()[:10])

兩個詞典中的部分 token 展示:

['', '[UNK]', '[START]', '[END]', '.', 'que', 'de', 'el', 'a', 'no']
['', '[UNK]', '[START]', '[END]', '.', 'the', 'i', 'to', 'you', 'tom']

(3)這裏主要是隨機選取一個樣本,我們將西班牙文中的 token 進行整數映射,然後再轉回西班牙文,很明顯我們看到,在句子的開頭和末尾加入了 [START] 、[END] ,而且如果在詞典中不存在的 token 直接表示為了 [UNK] 。英文的轉換過程也是如此。

def context_target_example(s, example, processor):
    print(s, ":")
    tokens = processor(example)
    print(tokens[:1, :])
    vocab = np.array(processor.get_vocabulary())
    result = vocab[tokens[0].numpy()]
    print(' '.join(result))

context_target_example('context', example_context ,context_processor )
context_target_example('target', example_target ,target_processor )

結果打印:

context :
<tf.RaggedTensor [[2, 8, 26, 646, 9, 204, 63, 117, 1, 4, 3]]>
[START] a los gatos no les gusta estar [UNK] . [END]
target :
<tf.RaggedTensor [[2, 677, 2399, 286, 1329, 4, 3]]>
[START] cats dislike being wet . [END]

(4)由於在 Seq2Seq 模型在 Decoder 部分,需要將目標進行預測,所以我們要對目標數據進行處理,使得每個時間步都有輸入和輸出,輸入當然就是其本身,而輸入就是相鄰的下一個 token 。我們在處理目標數據的時候,要將 target 數據本身當做 Decoder 輸入,然後整體將 target 右移一位當做 Decoder 輸出。

(5)我們隨機抽取了一個樣本,這裏展現了模型的 Encoder 輸入,Decoder 輸入以及 Docoder 輸出。我們可以看出在使用了 to_tensor 函數之後發生了填充操作,該 batch 中所有的 context 的長度都會變成該 batch 中出現的最長序列的那個 context 的長度,需要填充的位置都用 0 表示。同樣的道理該 batch 中的 target_in、target_out 的長度也都會發生同樣的填充變化。換句話説不同的 batch 中 Encoder 都不相等,不同的 batch 中 Decoder 的長度都不相等,唯一能保證相等的是同一個 batch 中的 Decoder 端的 target_in 和 target_out 長度相等。

def process_text(context, target):
    context = context_processor(context).to_tensor()
    target = target_processor(target)
    target_in = target[:,:-1].to_tensor()
    target_out = target[:,1:].to_tensor()
    return (context, target_in), target_out

train_datas = train_raw.map(process_text, tf.data.AUTOTUNE)
val_datas = val_raw.map(process_text, tf.data.AUTOTUNE)
for (one_context, one_target_in), one_target_out in train_datas.take(1):
    print("one_context:")
    print(one_context[:2, :].numpy()) 
    print("one_target_in:")
    print(one_target_in[:2, :].numpy()) 
    print("one_target_out:")
    print(one_target_out[:2, :].numpy())

文本處理後轉化為 token 的樣例:

one_context:
[[   2   13  627  616   14   20  610    9  605  134   12    3    0    0
     0    0    0    0]
 [   2   18   75  894    6   23  595    5   18 4656    4    3    0    0
     0    0    0    0]]
one_target_in:
[[  2  58 128  15   5 698  34  22 989  20   8  38  43  11]
 [  2  85 275  42  23  14  10 177  16   1  23   4   0   0]]
one_target_out:
[[ 58 128  15   5 698  34  22 989  20   8  38  43  11   3]
 [ 85 275  42  23  14  10 177  16   1  23   4   3   0   0]]

3. 搭建 Encoder

(1)第一層是文本處理器,要使用 context_processor 將文本進行預處理,並將 token 映射成整數 id

(2)第二層是嵌入層,要將每個整數 id 都映射成一個 128 維的向量。

(3)第三層是雙向 GRU 層,主要是捕獲輸入的西班牙文的文本特徵,並且在序列的每個時間步都輸出一個 128 維的向量

(4)我們用上面用過的例子 one_context 輸入到模型的 Encoder 中,可以看到輸入的大小為 [batch_size, source_seq_length] ,也就是説該 batch 中有 batch_size 個樣本,每個輸入樣本長度為 source_seq_length ,輸出的大小為 [batch_size, source_seq_length, units],也就是説該 batch 中有 batch_size 個樣本,每個輸入樣本的長度為 source_seq_length ,序列中每個時間步的輸出結果是 units 維。

UNITS = 128
class Encoder(tf.keras.layers.Layer):
    def __init__(self, text_processor, units):
        super(Encoder, self).__init__()
        self.text_processor = text_processor
        self.vocab_size = text_processor.vocabulary_size()
        self.units = units
        self.embedding = tf.keras.layers.Embedding(self.vocab_size, units, mask_zero=True)
        self.rnn = tf.keras.layers.Bidirectional( merge_mode='sum', layer=tf.keras.layers.GRU(units, return_sequences=True, recurrent_initializer='glorot_uniform'))

    def call(self, x):
        x = self.embedding(x)
        x = self.rnn(x)
        return x

    def convert_input(self, texts):
        texts = tf.convert_to_tensor(texts)
        if len(texts.shape) == 0:
            texts = tf.convert_to_tensor(texts)[tf.newaxis]
        context = self.text_processor(texts).to_tensor()
        context = self(context)
        return context

encoder = Encoder(context_processor, UNITS)
example_context_output = encoder(one_context)
print(f'Context tokens, shape (batch_size, source_seq_length)       : {one_context.shape}')
print(f'Encoder output, shape (batch_size, source_seq_length, units): {example_context_output.shape}')

結果輸出:

Context tokens, shape (batch_size, source_seq_length)       : (128, 18)
Encoder output, shape (batch_size, source_seq_length, units): (128, 18, 128)

4. 搭建 Attention

(1)Attention 層允許 Decoder 訪問 Encoder 提取的輸入文本的特徵信息,Attention 層會以 Decoder 輸出為 query ,以 Encoder 輸出為 key 和 value ,計算 Decoder 輸出與 Encoder 輸出的不同位置的有關重要程度,並將其加到 Decoder 的輸出中。

(2)我這裏選用了上面的例子 one_target_in ,假如它只經過了詞嵌入,然後輸出詞嵌入的結果,我們目前認為這就是經過 Decoder 解碼的輸出 example_target_embed ,我們計算這個 example_target_embed 和之前計算出來的對應的 Encoder 的輸出 example_context_output 的注意力結果向量,我們可以發現 Decoder 的輸入是 [batch_size, target_seq_length, units] ,表示 Decoder 的輸入的 batch 有 batch_size 個樣本,每個樣本長度為 target_seq_length,Decoder 的每個時間步上輸出的維度為 units 。 Attention 的結果和 Decoder 的向量維度是一樣的,這保證了 Attention 結果可以和 Decoder 輸出結果可以相加。Attention 的權重大小是 [batch_size, target_seq_length, source_seq_length] ,表示該 batch 中有 batch_size 個樣本,每個樣本的 Attention 的大小是 [target_seq_length, source_seq_length] ,表示計算出的每個樣本的 Decoder 輸出和 Encoder 輸出的所有一一對應的位置的有關重要程度。

class CrossAttention(tf.keras.layers.Layer):
    def __init__(self, units, **kwargs):
        super().__init__()
        self.mha = tf.keras.layers.MultiHeadAttention(key_dim=units, num_heads=1, **kwargs)
        self.layernorm = tf.keras.layers.LayerNormalization()
        self.add = tf.keras.layers.Add()

    def call(self, x, context):
        attn_output, attn_scores = self.mha( query=x,  value=context, return_attention_scores=True)
        attn_scores = tf.reduce_mean(attn_scores, axis=1)
        self.last_attention_weights = attn_scores
        x = self.add([x, attn_output])
        x = self.layernorm(x)
        return x

attention_layer = CrossAttention(UNITS)
embed = tf.keras.layers.Embedding(target_processor.vocabulary_size(),  output_dim=UNITS, mask_zero=True)
example_target_embed = embed(one_target_in)
result = attention_layer(example_target_embed, example_context_output)

print(f'Context sequence, shape (batch_size, source_seq_length, units)             : {example_context_output.shape}')
print(f'Target sequence, shape (batch_size, target_seq_length, units)              : {example_target_embed.shape}')
print(f'Attention result, shape (batch_size, target_seq_length, units)             : {result.shape}')
print(f'Attention weights, shape (batch_size, target_seq_length, source_seq_length): {attention_layer.last_attention_weights.shape}')

結果打印:

Context sequence, shape (batch_size, source_seq_length, units)             : (128, 18, 128)
Target sequence, shape (batch_size, target_seq_length, units)              : (128, 14, 128)
Attention result, shape (batch_size, target_seq_length, units)             : (128, 14, 128)
Attention weights, shape (batch_size, target_seq_length, source_seq_length): (128, 14, 18)

5. 搭建 Decoder

(1)第一層是文本處理器,使用 target_processor 將文本進行預處理,並將 token 映射成整數 id 。

(2)第二層是嵌入層,要將每個整數 id 都映射成一個 128 維的向量。

(3)第三層是一個單向 GRU 層,因為這裏是要從左到右進行的解碼工作,所以只能是一個從左到右的單向 GRU 層,主要是捕獲輸入的英文的文本特徵,並且在序列的每個時間步都輸出一個 128 維的向量。

(4)第四層是一個 Attention 層,這裏主要是計算第三層的輸出和 Encoder 的輸出的注意力結果,並將其和第三層的輸出進行相加。

(5)第五層是一個全連接層,Decoder 的每個輸出位置都有一個詞典大小的向量,表示每個位置預測下一個單詞的概率分佈。

(6)我們使用上面的例子產生的 Encoder 輸出 example_context_output 和 Decoder 輸入 one_target_in ,經過 Decoder 中間的計算過程,我們可以發現最終的輸出結果大小是 [batch_size, target_seq_length, target_vocabulary_size] 。表示輸出有 batch_size 個樣本結果,每個樣本的序列長度為 target_seq_length ,序列的每個位置上有一個詞典大小為 target_vocabulary_size 的概率分佈。

(7)在訓練時,模型會預測每個位置的目標單詞,每個位置的輸出預測結果都是沒有交互、獨立的,所以 Decoder 使用單向 GRU 來處理目標序列。但是使用模型進行推理時,每個位置生成一個單詞,並還要將此預測的單詞繼續反饋到模型的下一個位置中當作一部分輸入,進行下一個位置的預測。

class Decoder(tf.keras.layers.Layer):
    @classmethod
    def add_method(cls, fun):
        setattr(cls, fun.__name__, fun)
        return fun

    def __init__(self, text_processor, units):
        super(Decoder, self).__init__()
        self.text_processor = text_processor
        self.vocab_size = text_processor.vocabulary_size()
        self.word_to_id = tf.keras.layers.StringLookup( vocabulary=text_processor.get_vocabulary(),  mask_token='', oov_token='[UNK]')
        self.id_to_word = tf.keras.layers.StringLookup( vocabulary=text_processor.get_vocabulary(),  mask_token='', oov_token='[UNK]', invert=True)
        self.start_token = self.word_to_id('[START]')
        self.end_token = self.word_to_id('[END]')

        self.units = units
        self.embedding = tf.keras.layers.Embedding(self.vocab_size, units, mask_zero=True)
        self.rnn = tf.keras.layers.GRU(units, return_sequences=True, return_state=True,  recurrent_initializer='glorot_uniform')
        self.attention = CrossAttention(units)
        self.output_layer = tf.keras.layers.Dense(self.vocab_size)

    def call(self, context, x, state=None, return_state=False):  
        x = self.embedding(x)
        x, state = self.rnn(x, initial_state=state)
        x = self.attention(x, context)
        self.last_attention_weights = self.attention.last_attention_weights
        logits = self.output_layer(x)
        if return_state:
            return logits, state
        else:
            return logits

decoder = Decoder(target_processor, UNITS)
logits = decoder(example_context_output, one_target_in)

print(f'encoder output shape: (batch_size, source_seq_length, units) {example_context_output.shape}')
print(f'input target tokens shape: (batch_size, target_seq_length) {one_target_in.shape}')
print(f'logits shape shape: (batch_size, target_seq_length, target_vocabulary_size) {logits.shape}')

結果打印:

encoder output shape: (batch_size, source_seq_length, units) (128, 18, 128)
input target tokens shape: (batch_size, target_seq_length) (128, 14)
logits shape shape: (batch_size, target_seq_length, target_vocabulary_size) (128, 14, 6000)

(8)Decoder 在推理需要的其他需要的函數。

@Decoder.add_method
def get_initial_state(self, context):
    batch_size = tf.shape(context)[0]
    start_tokens = tf.fill([batch_size, 1], self.start_token)
    done = tf.zeros([batch_size, 1], dtype=tf.bool)
    embedded = self.embedding(start_tokens)
    return start_tokens, done, self.rnn.get_initial_state(embedded)[0]

@Decoder.add_method
def tokens_to_text(self, tokens):
    words = self.id_to_word(tokens)
    result = tf.strings.reduce_join(words, axis=-1, separator=' ')
    result = tf.strings.regex_replace(result, '^ *\[START\] *', '')
    result = tf.strings.regex_replace(result, ' *\[END\] *$', '')
    return result

@Decoder.add_method
def get_next_token(self, context, next_token, done, state, temperature = 0.0):
    logits, state = self( context, next_token, state = state, return_state=True) 
    if temperature == 0.0:
        next_token = tf.argmax(logits, axis=-1)
    else:
        logits = logits[:, -1, :]/temperature
        next_token = tf.random.categorical(logits, num_samples=1)
    done = done | (next_token == self.end_token)
    next_token = tf.where(done, tf.constant(0, dtype=tf.int64), next_token)
    return next_token, done, state

6. 搭建完整的 Translator 模型

(1)這裏創建翻譯類 Translator ,相當於一個完成 Seq2Seq 的模型,包含了一個 Encoder 和一個 Decoder 。

(2)使用之前用到的測試樣本,我們發現 Translator 返回的結果就是 Decoder 部分返回的結果。

class Translator(tf.keras.Model):
    @classmethod
    def add_method(cls, fun):
        setattr(cls, fun.__name__, fun)
        return fun

    def __init__(self, units,  context_processor, target_processor):
        super().__init__()
        encoder = Encoder(context_processor, units)
        decoder = Decoder(target_processor, units)
        self.encoder = encoder
        self.decoder = decoder

    def call(self, inputs):
        context, x = inputs
        context = self.encoder(context)
        logits = self.decoder(context, x)
        try:
            del logits._keras_mask
        except AttributeError:
            pass

        return logits

model = Translator(UNITS, context_processor, target_processor)
logits = model((one_context, one_target_in))

print(f'Context tokens, shape: (batch_size, source_seq_length) {one_context.shape}')
print(f'Target tokens, shape: (batch_size, target_seq_length) {one_target_in.shape}')
print(f'logits, shape: (batch_size, target_seq_length, target_vocabulary_size) {logits.shape}')

結果打印:

Context tokens, shape: (batch_size, source_seq_length) (128, 18)
Target tokens, shape: (batch_size, target_seq_length) (128, 14)
logits, shape: (batch_size, target_seq_length, target_vocabulary_size) (128, 14, 6000)

7. 編譯、訓練

(1)我們在這裏定義了模型使用了掩碼的損失函數的計算方法和準確率的計算方法,我們還選擇了 Adam 作為我們的優化器。

(2)在訓練過程中我們使用訓練數據訓練 100 個 epoch ,每個 epoch 訓練 100 個 batch ,並且在每個 epoch 訓練結束後使用驗證集進行評估,在第 24 個 epoch 的時候發生了 EarlyStopping 。

def masked_loss(y_true, y_pred):
    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy( from_logits=True, reduction='none')
    loss = loss_fn(y_true, y_pred)
    mask = tf.cast(y_true != 0, loss.dtype)
    loss *= mask
    return tf.reduce_sum(loss)/tf.reduce_sum(mask)

def masked_acc(y_true, y_pred):
    y_pred = tf.argmax(y_pred, axis=-1)
    y_pred = tf.cast(y_pred, y_true.dtype)
    match = tf.cast(y_true == y_pred, tf.float32)
    mask = tf.cast(y_true != 0, tf.float32)
    return tf.reduce_sum(match)/tf.reduce_sum(mask)

model.compile(optimizer='adam', loss=masked_loss,   metrics=[masked_acc, masked_loss])
history = model.fit( train_datas.repeat(),  epochs=100, steps_per_epoch = 100,  validation_data=val_datas, \
                    validation_steps = 20, callbacks=[ tf.keras.callbacks.EarlyStopping(patience=3)])

訓練過程:

Epoch 1/100
100/100 [==============================] - 43s 348ms/step - loss: 5.5715 - masked_acc: 0.2088 - masked_loss: 5.5715 - val_loss: 4.4883 - val_masked_acc: 0.3005 - val_masked_loss: 4.4883
Epoch 2/100
100/100 [==============================] - 31s 306ms/step - loss: 4.0577 - masked_acc: 0.3545 - masked_loss: 4.0577 - val_loss: 3.6292 - val_masked_acc: 0.4000 - val_masked_loss: 3.6292
......
Epoch 24/100
100/100 [==============================] - 32s 324ms/step - loss: 0.9762 - masked_acc: 0.7857 - masked_loss: 0.9762 - val_loss: 1.2726 - val_masked_acc: 0.7473 - val_masked_loss: 1.2726

(3)損失函數隨着 epoch 的變化過程。

plt.plot(history.history['loss'], label='loss')
plt.plot(history.history['val_loss'], label='val_loss')
plt.ylim([0, max(plt.ylim())])
plt.xlabel('Epoch')
plt.ylabel('loss')
plt.legend()

翻譯任務損失函數變化.png

(4)準確率隨着 epoch 的變化過程。

plt.plot(history.history['masked_acc'], label='accuracy')
plt.plot(history.history['val_masked_acc'], label='val_accuracy')
plt.ylim([0, max(plt.ylim())])
plt.xlabel('Epoch')
plt.ylabel('acc')
plt.legend()

翻譯任務準確率變化.png

8. 推理

(1)當模型訓練完成之後我們可以用訓練好的模型進行翻譯工作,這裏我們定義了一個翻譯函數,讓模型的翻譯的結果的最長長度為 50 。

@Translator.add_method
def translate(self, texts, *, max_length=500, temperature=tf.constant(0.0)):
    context = self.encoder.convert_input(texts)
    batch_size = tf.shape(context)[0]
    next_token, done, state = self.decoder.get_initial_state(context)
    tokens = tf.TensorArray(tf.int64, size=1, dynamic_size=True)
    for t in tf.range(max_length):
        next_token, done, state = self.decoder.get_next_token(  context, next_token, done, state, temperature)
        tokens = tokens.write(t, next_token)
        if tf.reduce_all(done):
            break
    tokens = tokens.stack()
    tokens = einops.rearrange(tokens, 't batch 1 -> batch t')
    text = self.decoder.tokens_to_text(tokens)
    return text

(2)我們使用了三個西班牙文的樣本來進行翻譯測試,並且我們給出了翻譯的耗時,可以看出翻譯基本準確。

%%time
inputs = [ 'Hace mucho frio aqui.',    # "It's really cold here."
           'Esta es mi vida.',         # "This is my life."
           'Su cuarto es un desastre.' # "His room is a mess"
         ]
for t in inputs:
    print(model.translate([t])[0].numpy().decode())

結果打印:

its very cold here . 
this is my life . 
his room is a disaster . 
CPU times: total: 577 ms
Wall time: 578 ms

9. 保存和加載模型

(1)我們要將訓練好的模型進行保存,在使用的時候可以進行加載使用。

class Export(tf.Module):
    def __init__(self, model):
        self.model = model
    @tf.function(input_signature=[tf.TensorSpec(dtype=tf.string, shape=[None])])
    def translate(self, inputs):
        return self.model.translate(inputs)
export = Export(model)
tf.saved_model.save(export, 'dynamic_translator',  signatures={'serving_default': export.translate})

(2)使用加載的模型進行推理

%%time
reloaded = tf.saved_model.load('dynamic_translator')
result = reloaded.translate(tf.constant(inputs))
print(result[0].numpy().decode())
print(result[1].numpy().decode())
print(result[2].numpy().decode())

結果打印:

its very cold here .  
this is my life .  
his room is a disaster . 
CPU times: total: 42.5 s
Wall time: 42.8 s