备注
点击 这里 下载完整示例代码
使用 TorchScript 部署 Seq2Seq 模型¶
Created On: Sep 17, 2018 | Last Updated: Dec 02, 2024 | Last Verified: Nov 05, 2024
警告
TorchScript 不再处于活跃开发中。
本教程将逐步讲解如何使用 TorchScript API 将序列到序列模型迁移到 TorchScript。我们将转换的模型是来自 Chatbot 教程 的聊天机器人模型。您可以将此教程视为 Chatbot 教程的“第二部分”并部署您自己的预训练模型，或者您可以从本文档开始，并使用我们托管的预训练模型。在后一种情况下，您可以参考原始 Chatbot 教程，以了解数据预处理、模型理论及定义和模型训练的详细信息。
什么是 TorchScript？¶
在基于深度学习的项目的研究和开发阶段，使用像 PyTorch 这样直接而命令式的**实时**界面非常有利。这使用户能够编写熟悉的、符合 Python 规范的代码，可以使用 Python 数据结构、控制流操作、打印语句和调试工具。虽然实时界面对研究和实验有益，但当需要在生产环境中部署模型时，具有**图**形式的模型表示形式非常有用。推迟的图形式表示可以实现优化，例如乱序执行，并能够适配高效的硬件架构。此外，基于图的表示形式还可以实现框架无关的模型导出。PyTorch 提供了将实时模式代码逐步转换为 TorchScript 的机制，TorchScript 是独立于 Python 运行时的可静态分析和优化的 Python 子集，Torch 使用它来表示深度学习程序。
将实时模式的 PyTorch 程序转换为 TorchScript 的 API 位于 torch.jit 模块中。该模块为将实时模式模型转换为 TorchScript 图形表示提供了两种核心模式：**跟踪**（Tracing）和**脚本**（Scripting）。torch.jit.trace 函数接受一个模块或函数以及一组示例输入。然后，它会通过函数或模块运行示例输入，同时记录遇到的计算步骤，并输出一个基于图形的函数，用于执行记录的操作。**跟踪**非常适合那些没有数据依赖控制流的常规模块和函数，例如标准卷积神经网络。然而，如果对包含数据依赖的 if 语句和循环的函数进行跟踪，那么只有示例输入所采取的执行路径上的操作会被记录。换句话说，控制流本身不会被捕获。为了转换包含数据依赖控制流的模块和函数，还提供了一种**脚本**机制。torch.jit.script 函数/装饰器接受一个模块或函数，不需要示例输入。脚本会显式地将模块或函数代码转换为 TorchScript，包括所有控制流。脚本使用的一个警告是它仅支持 Python 的一个子集，因此您可能需要重写代码以使其兼容 TorchScript 语法。
有关所有支持功能的详细信息，请参阅 TorchScript 语言参考。为了提供最大灵活性，您还可以混合使用跟踪和脚本模式来表示整个程序，可以渐进式地应用这些技术。
归功与致谢¶
本教程受以下来源的启发：
Yuan-Kuei Wu 的 pytorch-chatbot 实现: https://github.com/ywk991112/pytorch-chatbot
Sean Robertson 的 practical-pytorch 序列到序列翻译示例: https://github.com/spro/practical-pytorch/tree/master/seq2seq-translation
FloydHub 的 Cornell Movie Corpus 数据预处理代码: https://github.com/floydhub/textutil-preprocess-cornell-movie-corpus
准备环境¶
首先，我们将导入所需的模块并设置一些常量。如果您打算使用自己的模型，请确保 MAX_LENGTH 常量设置正确。请记住，这个常量定义了训练期间允许的最大句子长度，以及模型能够生成的最大输出长度。
import torch
import torch.nn as nn
import torch.nn.functional as F
import re
import os
import unicodedata
import numpy as np
device = torch.device("cpu")
MAX_LENGTH = 10 # Maximum sentence length
# Default word tokens
PAD_token = 0 # Used for padding short sentences
SOS_token = 1 # Start-of-sentence token
EOS_token = 2 # End-of-sentence token
模型概述¶
如前所述，我们使用的模型是一个 `序列到序列 <https://arxiv.org/abs/1409.3215>`__（seq2seq）模型。这种类型的模型用于输入是可变长度序列，并且输出也是可变的长度序列，且不一定是一对一的映射。一个 seq2seq 模型包含两个循环神经网络（RNN），它们协同工作：一个 编码器 和一个 解码器。
图片来源：https://jeddy92.github.io/JEddy92.github.io/ts_seq2seq_intro/
编码器¶
编码器 RNN 通过输入句子逐个标记（例如：单词）进行迭代，在每个时间步输出一个“输出”向量和一个“隐藏状态”向量。隐藏状态向量随后传递到下一时间步，而输出向量被记录下来。编码器将它在序列中各点看到的上下文转换为一个高维空间中的一组点，解码器将使用这些点来为给定任务生成有意义的输出。
数据处理¶
尽管我们的模型在概念上处理的是令牌序列，但实际中，它们像所有机器学习模型一样处理数字。在这种情况下，模型的词汇中的每个单词（在训练之前被建立）都映射到一个整数索引中。我们使用一个 Voc 对象来包含从单词到索引的映射以及词汇表中的总词数。稍后在运行模型之前我们会加载该对象。
而且，为了能够运行评估，我们必须提供一个工具来处理我们的字符串输入。normalizeString 函数将字符串中的所有字符转换为小写并移除所有非字母字符。indexesFromSentence 函数接收一个由单词组成的句子并返回对应的单词索引序列。
class Voc:
def __init__(self, name):
self.name = name
self.trimmed = False
self.word2index = {}
self.word2count = {}
self.index2word = {PAD_token: "PAD", SOS_token: "SOS", EOS_token: "EOS"}
self.num_words = 3 # Count SOS, EOS, PAD
def addSentence(self, sentence):
for word in sentence.split(' '):
self.addWord(word)
def addWord(self, word):
if word not in self.word2index:
self.word2index[word] = self.num_words
self.word2count[word] = 1
self.index2word[self.num_words] = word
self.num_words += 1
else:
self.word2count[word] += 1
# Remove words below a certain count threshold
def trim(self, min_count):
if self.trimmed:
return
self.trimmed = True
keep_words = []
for k, v in self.word2count.items():
if v >= min_count:
keep_words.append(k)
print('keep_words {} / {} = {:.4f}'.format(
len(keep_words), len(self.word2index), len(keep_words) / len(self.word2index)
))
# Reinitialize dictionaries
self.word2index = {}
self.word2count = {}
self.index2word = {PAD_token: "PAD", SOS_token: "SOS", EOS_token: "EOS"}
self.num_words = 3 # Count default tokens
for word in keep_words:
self.addWord(word)
# Lowercase and remove non-letter characters
def normalizeString(s):
s = s.lower()
s = re.sub(r"([.!?])", r" \1", s)
s = re.sub(r"[^a-zA-Z.!?]+", r" ", s)
return s
# Takes string sentence, returns sentence of word indexes
def indexesFromSentence(voc, sentence):
return [voc.word2index[word] for word in sentence.split(' ')] + [EOS_token]
定义编码器¶
我们使用 torch.nn.GRU 模块实现了编码器的 RNN，该模块接收一个句子批次（单词嵌入的向量），并内部按顺序逐个标记计算隐藏状态。我们将该模块初始化为双向模式，这意味着我们有两个独立的 GRU：一个按照时间顺序迭代序列，另一个按照反向顺序迭代序列。最终，我们返回这两个 GRU 输出的总和。由于我们的模型是通过批处理训练的，因此 EncoderRNN 模型的 forward 方法期望接收到一个填充过的输入批次。为了批量化可变长度的句子，我们允许句子中最多包含 MAX_LENGTH 个标记，所有少于 MAX_LENGTH 个标记的句子将在结尾处用我们专门的 PAD_token 标记进行填充。为了在使用 PyTorch RNN 模块时处理填充批次，我们需要使用 torch.nn.utils.rnn.pack_padded_sequence 和 torch.nn.utils.rnn.pad_packed_sequence 数据转换包装前向传递调用。注意，forward 方法还接受一个 input_lengths 列表，其中包含批次中每个句子的长度。此输入由 torch.nn.utils.rnn.pack_padded_sequence 函数在进行填充时使用。
TorchScript 注意事项：¶
由于编码器的 forward 方法不包含任何数据相关的控制流，因此我们将使用 tracing 来将其转换为脚本模式。在跟踪模块时，我们可以保持模块定义不变。我们将在本文档的末尾初始化所有模型，然后再运行评估。
class EncoderRNN(nn.Module):
def __init__(self, hidden_size, embedding, n_layers=1, dropout=0):
super(EncoderRNN, self).__init__()
self.n_layers = n_layers
self.hidden_size = hidden_size
self.embedding = embedding
# Initialize GRU; the ``input_size`` and ``hidden_size`` parameters are both set to 'hidden_size'
# because our input size is a word embedding with number of features == hidden_size
self.gru = nn.GRU(hidden_size, hidden_size, n_layers,
dropout=(0 if n_layers == 1 else dropout), bidirectional=True)
def forward(self, input_seq, input_lengths, hidden=None):
# type: (Tensor, Tensor, Optional[Tensor]) -> Tuple[Tensor, Tensor]
# Convert word indexes to embeddings
embedded = self.embedding(input_seq)
# Pack padded batch of sequences for RNN module
packed = torch.nn.utils.rnn.pack_padded_sequence(embedded, input_lengths)
# Forward pass through GRU
outputs, hidden = self.gru(packed, hidden)
# Unpack padding
outputs, _ = torch.nn.utils.rnn.pad_packed_sequence(outputs)
# Sum bidirectional GRU outputs
outputs = outputs[:, :, :self.hidden_size] + outputs[:, : ,self.hidden_size:]
# Return output and final hidden state
return outputs, hidden
定义解码器的注意力模块¶
接下来，我们将定义注意力模块（Attn）。注意，此模块将用作解码器模型中的子模块。Luong 等人考虑了各种“评分函数”，这些函数会接收当前解码器 RNN 的输出和整个编码器的输出，并返回注意力“能量”。此注意力能量张量与编码器输出的大小相同，两者最终相乘，结果是一个加权张量，其最大值表示查询句子在解码的特定时间步中最重要的部分。
# Luong attention layer
class Attn(nn.Module):
def __init__(self, method, hidden_size):
super(Attn, self).__init__()
self.method = method
if self.method not in ['dot', 'general', 'concat']:
raise ValueError(self.method, "is not an appropriate attention method.")
self.hidden_size = hidden_size
if self.method == 'general':
self.attn = nn.Linear(self.hidden_size, hidden_size)
elif self.method == 'concat':
self.attn = nn.Linear(self.hidden_size * 2, hidden_size)
self.v = nn.Parameter(torch.FloatTensor(hidden_size))
def dot_score(self, hidden, encoder_output):
return torch.sum(hidden * encoder_output, dim=2)
def general_score(self, hidden, encoder_output):
energy = self.attn(encoder_output)
return torch.sum(hidden * energy, dim=2)
def concat_score(self, hidden, encoder_output):
energy = self.attn(torch.cat((hidden.expand(encoder_output.size(0), -1, -1), encoder_output), 2)).tanh()
return torch.sum(self.v * energy, dim=2)
def forward(self, hidden, encoder_outputs):
# Calculate the attention weights (energies) based on the given method
if self.method == 'general':
attn_energies = self.general_score(hidden, encoder_outputs)
elif self.method == 'concat':
attn_energies = self.concat_score(hidden, encoder_outputs)
elif self.method == 'dot':
attn_energies = self.dot_score(hidden, encoder_outputs)
# Transpose max_length and batch_size dimensions
attn_energies = attn_energies.t()
# Return the softmax normalized probability scores (with added dimension)
return F.softmax(attn_energies, dim=1).unsqueeze(1)
定义解码器¶
与 EncoderRNN 类似，我们使用 torch.nn.GRU 模块作为解码器的 RNN。然而，这次我们只使用单向 GRU。需要注意的是，不同于编码器，我们会一次喂给解码器 RNN 一个单词。首先我们获取当前单词的嵌入，并应用 dropout。接下来，我们将嵌入和最后的隐藏状态传递给 GRU，并获取当前的 GRU 输出和隐藏状态。然后我们使用 Attn 模块作为一层来获得注意力权重，这些权重与编码器的输出相乘以获得我们关注的编码器输出。我们使用此关注的编码器输出作为 context 张量，该张量表示一个加权和，用于指示编码器输出的哪些部分需要关注。从这里，我们使用线性层和 softmax 正规化来选择输出序列中的下一个单词。
# TorchScript Notes:
# ~~~~~~~~~~~~~~~~~~~~~~
#
# Similarly to the ``EncoderRNN``, this module does not contain any
# data-dependent control flow. Therefore, we can once again use
# **tracing** to convert this model to TorchScript after it
# is initialized and its parameters are loaded.
#
class LuongAttnDecoderRNN(nn.Module):
def __init__(self, attn_model, embedding, hidden_size, output_size, n_layers=1, dropout=0.1):
super(LuongAttnDecoderRNN, self).__init__()
# Keep for reference
self.attn_model = attn_model
self.hidden_size = hidden_size
self.output_size = output_size
self.n_layers = n_layers
self.dropout = dropout
# Define layers
self.embedding = embedding
self.embedding_dropout = nn.Dropout(dropout)
self.gru = nn.GRU(hidden_size, hidden_size, n_layers, dropout=(0 if n_layers == 1 else dropout))
self.concat = nn.Linear(hidden_size * 2, hidden_size)
self.out = nn.Linear(hidden_size, output_size)
self.attn = Attn(attn_model, hidden_size)
def forward(self, input_step, last_hidden, encoder_outputs):
# Note: we run this one step (word) at a time
# Get embedding of current input word
embedded = self.embedding(input_step)
embedded = self.embedding_dropout(embedded)
# Forward through unidirectional GRU
rnn_output, hidden = self.gru(embedded, last_hidden)
# Calculate attention weights from the current GRU output
attn_weights = self.attn(rnn_output, encoder_outputs)
# Multiply attention weights to encoder outputs to get new "weighted sum" context vector
context = attn_weights.bmm(encoder_outputs.transpose(0, 1))
# Concatenate weighted context vector and GRU output using Luong eq. 5
rnn_output = rnn_output.squeeze(0)
context = context.squeeze(1)
concat_input = torch.cat((rnn_output, context), 1)
concat_output = torch.tanh(self.concat(concat_input))
# Predict next word using Luong eq. 6
output = self.out(concat_output)
output = F.softmax(output, dim=1)
# Return output and final hidden state
return output, hidden
定义评估¶
贪婪搜索解码器¶
与聊天机器人教程一样，我们使用 GreedySearchDecoder 模块来方便实际解码过程。此模块将训练的编码器和解码器模型作为属性，并驱动将输入句子（单词索引的向量）编码，以及迭代地解码输出响应序列的过程，一次一个单词（单词索引）。
编码输入序列很简单：只需将整个序列张量及其对应的长度向量传递到 encoder。需要注意此模块一次只处理一个输入序列，而不是**序列批次。因此，当常数 **1 用于声明张量大小时，它对应于批量大小为 1。要解码给定解码器的输出，我们必须迭代地通过解码器模型运行前向传递，该模型输出与每个单词是解码序列中正确的下一个单词的概率相对应的 softmax 分数。我们将 decoder_input 初始化为一个包含 SOS_token 的张量。每次通过解码器后，我们将具有最高 softmax 概率的单词“贪婪地”附加到 decoded_words 列表中。我们还将此单词用作下一次迭代的 decoder_input。解码过程的结束条件是 decoded_words 列表已达到 MAX_LENGTH 长度，或预测的单词为 EOS_token。
TorchScript 注意事项：¶
该模块的 forward 方法在解码输出序列时涉及对范围 \([0, max\_length)\) 的迭代，因为我们一次解码一个单词。因此，我们应该使用 scripting 将此模块转换为 TorchScript。不像我们的编码器和解码器模型，我们可以跟踪，在初始化对象时，我们需要对 GreedySearchDecoder 模块进行一些必要的更改，以确保不会报错。换句话说，我们必须确保模块遵守 TorchScript 机制的规则，并且不会使用任何超出 TorchScript 所支持的 Python 子集的语言功能。
为了了解一些可能需要的操作，我们将比较聊天机器人教程中的 GreedySearchDecoder 实现和我们在以下单元中使用的实现的异同。请注意，红色高亮的行是从原始实现中删除的行，绿色高亮的行是新添加的行。
更改：¶
向构造函数参数添加
decoder_n_layers此更改源于我们传递给此模块的编码器和解码器模型将是
TracedModule的子模块（而不是Module）。因此，我们无法通过decoder.n_layers来访问解码器的层数。相反，我们对此进行了规划，并在模块构建过程中传入此值。
将新的属性存储为常量
在原始实现中，我们可以自由使用
GreedySearchDecoder的forward方法中周围（全局）范围的变量。然而现在我们使用 scripting，因为 scripting 假设我们不能必然地保留 Python 对象，尤其是在导出时。一个简单的解决方案是将这些全局范围的值作为属性存储到模块中，并将它们添加到一个特殊列表__constants__中，以便在forward方法构造图时将其用作字面值。在 NEW 的第 19 行中，我们不仅使用了device和SOS_token的全局值，而且使用了作为常量属性的self._device和self._SOS_token。
强制规定
forward方法参数的类型默认情况下，所有 TorchScript 函数的参数都假定为 Tensor。如果需要传递不同类型的参数，可以使用 PEP 3107 中介绍的函数类型注解。此外，还可以使用 Mypy 样式的类型注解来声明不同类型的参数（详情见 文档）。
更改
decoder_input的初始化在原始实现中，我们通过
torch.LongTensor([[SOS_token]])初始化了decoder_input张量。当使用脚本时，不允许以字面方式初始化张量。相反，我们可以使用显式的 torch 函数（如torch.ones）来初始化张量。在这种情况下，我们可以通过将常量self._SOS_token的 SOS_token 值乘以 1，轻松复制标量decoder_input张量。
class GreedySearchDecoder(nn.Module):
def __init__(self, encoder, decoder, decoder_n_layers):
super(GreedySearchDecoder, self).__init__()
self.encoder = encoder
self.decoder = decoder
self._device = device
self._SOS_token = SOS_token
self._decoder_n_layers = decoder_n_layers
__constants__ = ['_device', '_SOS_token', '_decoder_n_layers']
def forward(self, input_seq : torch.Tensor, input_length : torch.Tensor, max_length : int):
# Forward input through encoder model
encoder_outputs, encoder_hidden = self.encoder(input_seq, input_length)
# Prepare encoder's final hidden layer to be first hidden input to the decoder
decoder_hidden = encoder_hidden[:self._decoder_n_layers]
# Initialize decoder input with SOS_token
decoder_input = torch.ones(1, 1, device=self._device, dtype=torch.long) * self._SOS_token
# Initialize tensors to append decoded words to
all_tokens = torch.zeros([0], device=self._device, dtype=torch.long)
all_scores = torch.zeros([0], device=self._device)
# Iteratively decode one word token at a time
for _ in range(max_length):
# Forward pass through decoder
decoder_output, decoder_hidden = self.decoder(decoder_input, decoder_hidden, encoder_outputs)
# Obtain most likely word token and its softmax score
decoder_scores, decoder_input = torch.max(decoder_output, dim=1)
# Record token and score
all_tokens = torch.cat((all_tokens, decoder_input), dim=0)
all_scores = torch.cat((all_scores, decoder_scores), dim=0)
# Prepare current token to be next decoder input (add a dimension)
decoder_input = torch.unsqueeze(decoder_input, 0)
# Return collections of word tokens and scores
return all_tokens, all_scores
评估输入¶
接下来，我们定义一些函数来评估输入。evaluate 函数接收一个规范化的字符串句子，将其处理为对应单词索引的张量（批量大小为 1），并将该张量传递给一个名为 searcher 的 GreedySearchDecoder 实例来处理编码/解码流程。搜索器返回输出单词索引向量和一个对应于每个解码单词标记的 softmax 分数张量。最后一步是使用 voc.index2word 将每个单词索引转换回其字符串表示形式。
我们还定义了两个函数来评估输入句子。evaluateInput 函数会提示用户输入内容并进行评估。它会不断要求输入，直到用户输入 ‘q’ 或 ‘quit’。
evaluateExample 函数只需要一个字符串输入句子作为参数，规范化后进行评估，并打印响应。
def evaluate(searcher, voc, sentence, max_length=MAX_LENGTH):
### Format input sentence as a batch
# words -> indexes
indexes_batch = [indexesFromSentence(voc, sentence)]
# Create lengths tensor
lengths = torch.tensor([len(indexes) for indexes in indexes_batch])
# Transpose dimensions of batch to match models' expectations
input_batch = torch.LongTensor(indexes_batch).transpose(0, 1)
# Use appropriate device
input_batch = input_batch.to(device)
lengths = lengths.to(device)
# Decode sentence with searcher
tokens, scores = searcher(input_batch, lengths, max_length)
# indexes -> words
decoded_words = [voc.index2word[token.item()] for token in tokens]
return decoded_words
# Evaluate inputs from user input (``stdin``)
def evaluateInput(searcher, voc):
input_sentence = ''
while(1):
try:
# Get input sentence
input_sentence = input('> ')
# Check if it is quit case
if input_sentence == 'q' or input_sentence == 'quit': break
# Normalize sentence
input_sentence = normalizeString(input_sentence)
# Evaluate sentence
output_words = evaluate(searcher, voc, input_sentence)
# Format and print response sentence
output_words[:] = [x for x in output_words if not (x == 'EOS' or x == 'PAD')]
print('Bot:', ' '.join(output_words))
except KeyError:
print("Error: Encountered unknown word.")
# Normalize input sentence and call ``evaluate()``
def evaluateExample(sentence, searcher, voc):
print("> " + sentence)
# Normalize sentence
input_sentence = normalizeString(sentence)
# Evaluate sentence
output_words = evaluate(searcher, voc, input_sentence)
output_words[:] = [x for x in output_words if not (x == 'EOS' or x == 'PAD')]
print('Bot:', ' '.join(output_words))
加载预训练参数¶
现在，让我们加载模型！
使用托管模型¶
要加载托管模型：
在 这里 下载模型。
将
loadFilename变量设置为下载的检查点文件的路径。保持
checkpoint = torch.load(loadFilename)行未注释，因为托管模型是在 CPU 上训练的。
使用您自己的模型¶
要加载您自己的预训练模型：
将
loadFilename变量设置为您希望加载的检查点文件的路径。注意，如果您遵循聊天机器人教程中保存模型的约定，这可能需要更改model_name、encoder_n_layers、decoder_n_layers、hidden_size和 ``checkpoint_iter``（因为这些值在模型路径中使用）。如果您在 CPU 上训练模型，请确保您使用
checkpoint = torch.load(loadFilename)行打开检查点。如果您在 GPU 上训练模型，而正在使用此教程的 CPU，请取消注释checkpoint = torch.load(loadFilename, map_location=torch.device('cpu'))行。
TorchScript 注意事项：¶
注意，我们像往常一样初始化并加载参数到我们的编码器和解码器模型中。如果您对某些模型部分使用了跟踪模式（torch.jit.trace），在跟踪模型之前必须调用 .to(device) 设置模型的设备选项，以及调用 .eval() 将 dropout 层设置为测试模式。TracedModule 对象不继承 to 或 eval 方法。由于在本教程中我们仅使用 scripting 而非 tracing，我们只需要在评估之前进行这些操作（与通常的 eager 模式相同）。
save_dir = os.path.join("data", "save")
corpus_name = "cornell movie-dialogs corpus"
# Configure models
model_name = 'cb_model'
attn_model = 'dot'
#attn_model = 'general'``
#attn_model = 'concat'
hidden_size = 500
encoder_n_layers = 2
decoder_n_layers = 2
dropout = 0.1
batch_size = 64
# If you're loading your own model
# Set checkpoint to load from
checkpoint_iter = 4000
从检查点加载的示例代码：
loadFilename = os.path.join(save_dir, model_name, corpus_name,
'{}-{}_{}'.format(encoder_n_layers, decoder_n_layers, hidden_size),
'{}_checkpoint.tar'.format(checkpoint_iter))
# If you're loading the hosted model
loadFilename = 'data/4000_checkpoint.tar'
# Load model
# Force CPU device options (to match tensors in this tutorial)
checkpoint = torch.load(loadFilename, map_location=torch.device('cpu'))
encoder_sd = checkpoint['en']
decoder_sd = checkpoint['de']
encoder_optimizer_sd = checkpoint['en_opt']
decoder_optimizer_sd = checkpoint['de_opt']
embedding_sd = checkpoint['embedding']
voc = Voc(corpus_name)
voc.__dict__ = checkpoint['voc_dict']
print('Building encoder and decoder ...')
# Initialize word embeddings
embedding = nn.Embedding(voc.num_words, hidden_size)
embedding.load_state_dict(embedding_sd)
# Initialize encoder & decoder models
encoder = EncoderRNN(hidden_size, embedding, encoder_n_layers, dropout)
decoder = LuongAttnDecoderRNN(attn_model, embedding, hidden_size, voc.num_words, decoder_n_layers, dropout)
# Load trained model parameters
encoder.load_state_dict(encoder_sd)
decoder.load_state_dict(decoder_sd)
# Use appropriate device
encoder = encoder.to(device)
decoder = decoder.to(device)
# Set dropout layers to ``eval`` mode
encoder.eval()
decoder.eval()
print('Models built and ready to go!')
将模型转换为 TorchScript¶
编码器¶
如前所述，要将编码器模型转换为 TorchScript，我们使用 scripting。编码器模型接收一个输入序列和一个对应的长度张量。因此，我们创建一个样例输入序列张量 test_seq，这个张量具有适当的大小（MAX_LENGTH, 1），包含适当范围 \([0, voc.num\_words)\) 的数字，且具有适当的类型（int64）。我们还创建了一个 test_seq_length 标量，该标量实际上包含与 test_seq 中单词数量对应的值。下一步是使用 torch.jit.trace 函数跟踪模型。注意我们传递的第一个参数是我们想要跟踪的模块，第二个参数是模块 forward 方法的参数元组。
解码器¶
我们对解码器进行的跟踪过程与对编码器进行的过程相同。注意，我们通过在 traced_encoder 上使用一组随机输入调用 forward 方法来获取解码器所需的输出。这不是必须的，因为我们同样可以简单地制造一个具有正确形状、类型和值范围的张量。该方法可行，因为在我们的情况下，张量的值没有任何限制，因为我们没有任何可能因输入值超出范围而出错的操作。
贪婪搜索解码器¶
记得我们因为存在数据依赖的控制流而对搜索模块进行了脚本化。在脚本化的情况下，我们需要做必要的语言更改，以确保实现符合 TorchScript 的要求。我们初始化脚本化的搜索器与初始化未脚本化的变体方式相同。
### Compile the whole greedy search model to TorchScript model
# Create artificial inputs
test_seq = torch.LongTensor(MAX_LENGTH, 1).random_(0, voc.num_words).to(device)
test_seq_length = torch.LongTensor([test_seq.size()[0]]).to(device)
# Trace the model
traced_encoder = torch.jit.trace(encoder, (test_seq, test_seq_length))
### Convert decoder model
# Create and generate artificial inputs
test_encoder_outputs, test_encoder_hidden = traced_encoder(test_seq, test_seq_length)
test_decoder_hidden = test_encoder_hidden[:decoder.n_layers]
test_decoder_input = torch.LongTensor(1, 1).random_(0, voc.num_words)
# Trace the model
traced_decoder = torch.jit.trace(decoder, (test_decoder_input, test_decoder_hidden, test_encoder_outputs))
### Initialize searcher module by wrapping ``torch.jit.script`` call
scripted_searcher = torch.jit.script(GreedySearchDecoder(traced_encoder, traced_decoder, decoder.n_layers))
打印图表¶
现在我们的模型已转换为 TorchScript 形式，我们可以打印每个模型的图表，以确保我们正确捕获了计算图。由于 TorchScript 允许我们递归编译整个模型层级，并将 encoder 和 decoder 的图嵌入到一个单一图中，我们只需要打印 scripted_searcher 的图表即可。
print('scripted_searcher graph:\n', scripted_searcher.graph)
运行评估¶
最后，我们将使用 TorchScript 模型运行聊天机器人模型的评估。如果转换正确，模型的行为将与它们在急迫模式的表示中完全一致。
默认情况下，我们会评估几个常见的查询语句。如果你想亲自与机器人聊天，可以取消注释 evaluateInput 行并尝试一下。
# Use appropriate device
scripted_searcher.to(device)
# Set dropout layers to ``eval`` mode
scripted_searcher.eval()
# Evaluate examples
sentences = ["hello", "what's up?", "who are you?", "where am I?", "where are you from?"]
for s in sentences:
evaluateExample(s, scripted_searcher, voc)
# Evaluate your input by running
# ``evaluateInput(traced_encoder, traced_decoder, scripted_searcher, voc)``
保存模型¶
现在我们已经成功将模型转换为 TorchScript，我们将对其进行序列化，以便在非 Python 的部署环境中使用。为此，我们可以简单地保存 scripted_searcher 模块，因为这是进行聊天机器人模型推理时的用户接口。在保存 Script 模块时，应使用 script_module.save(PATH) 而不是 torch.save(model, PATH)。
scripted_searcher.save("scripted_chatbot.pth")
脚本总运行时间： (0 分钟 0.000 秒)
