Note:
Click here to download the full example code
Transformer tutorial
Author: Zihao Ye, Jinjing Zhou, Qipeng Guo, Quan Gan, Zheng Zhang
在本教程中，您将学习Transformer模型的简化实现。您可以看到最重要的设计要点的亮点。例如，只有单头注意力。完整的代码可以在这里找到 。
总体结构与研究论文 Annotated Transformer中的结构相似。
在研究论文中介绍了Transformer模型，以替代CNN / RNN体系结构进行序列建模，这就是“ Attention is All You Need”。它改善了机器翻译以及自然语言推理任务（GPT）的技术水平。带有大型语料库（BERT）的预训练Transformer的最新工作支持它能够学习高质量的语义表示。
变形金刚有趣的部分是其广泛的关注度。注意的经典用法来自机器翻译模型，其中输出令牌出现在所有输入令牌上。
变压器还会在解码器和编码器中施加自注意。不论单词在序列中的位置如何，此过程都会迫使单词相互关联以组合在一起。这与基于RNN的模型不同，在RNN的模型中，单词（在源句子中）沿链组合在一起，这被认为过于局限。
Attention layer of Transformer
在Transformer的关注层中，模块会为每个节点学习在其传入边缘上分配权重。对于节点对(i,j) （从i 至 j）与节点xi,xj∈Rn，它们的连接分数定义如下：
qj=Wq⋅xjki=Wk⋅xivi=Wv⋅xiscore=qjTki
哪里 Wq,Wk,Wv∈Rn×dk 映射表示 x 分别为“查询”，“键”和“值”空间。
还有其他实现分数功能的可能性。点积可衡量给定查询的相似性qj 和一把钥匙 ki：如果 j 需要存储在其中的信息 i，位置的查询向量 j （qj）应该在位置处接近关键向量 i （ki）。
然后将分数用于计算存储在的权重上归一化的输入值的总和 wv。然后将仿射层应用于wv 得到输出 o：
wji=(k,i)∈E∑exp{scoreki}exp{scoreji}wvi=(k,i)∈E∑wkivko=Wo⋅wv
Multi-head attention layer
在《变形金刚》中，注意力是多头的。头非常像卷积网络中的通道。多头注意力由多个注意力头组成，其中每个头都指一个注意力模块。wv(i) 所有头都被串联并映射到输出 o 仿射层：
o=Wo⋅concat([wv(0),wv(1),⋯,wv(h)])
下面的代码包装了用于多头注意的必要组件，并提供了两个接口。
- get将状态“ x”映射到查询，键和值，这是后续步骤（propagate_attention）所必需的。
- get_o 将关注后的更新值映射到输出 o 用于后期处理。
class MultiHeadAttention(nn.Module):
"Multi-Head Attention"
def __init__(self, h, dim_model):
"h: number of heads; dim_model: hidden dimension"
super(MultiHeadAttention, self).__init__()
self.d_k = dim_model // h
self.h = h
# W_q, W_k, W_v, W_o
self.linears = clones(nn.Linear(dim_model, dim_model), 4)
def get(self, x, fields='qkv'):
"Return a dict of queries / keys / values."
batch_size = x.shape[0]
ret = {}
if 'q' in fields:
ret['q'] = self.linears[0](x).view(batch_size, self.h, self.d_k)
if 'k' in fields:
ret['k'] = self.linears[1](x).view(batch_size, self.h, self.d_k)
if 'v' in fields:
ret['v'] = self.linears[2](x).view(batch_size, self.h, self.d_k)
return ret
def get_o(self, x):
"get output of the multi-head attention"
batch_size = x.shape[0]
return self.linears[3](x.view(batch_size, -1))
How DGL implements Transformer with a graph neural network
通过将注意力视为图形中的边缘并采用在边缘上传递的消息来引发适当的处理，您将获得Transformer的不同视角。
Graph structure
通过将源句子和目标句子的标记映射到节点来构造图。完整的Transformer图由三个子图组成：
源语言图。这是一个完整的图形，每个标记si 可以参加任何其他令牌 sj（包括自循环）。
图片地址：https://i.imgur.com/zV5LmTX.png
目标语言图。该图是半完整的，因为ti 只参加 tj 如果 i>j（输出令牌不能取决于将来的单词）。
图片地址：https://i.imgur.com/dETQMMx.png
跨语言图。这是一个双向图，其中每个源令牌都有一条边si 每个目标代币 tj，这意味着每个目标令牌都可以参加源令牌。
图片地址：https://i.imgur.com/hnGP229.png
完整图片如下所示：
图片地址：https://i.imgur.com/Hj2rRGT.png
在数据集准备阶段预先构建图形。
Message passing
定义图结构后，继续定义用于消息传递的计算。
假设您已经计算了所有查询 qi，键 ki 和价值观 vi。对于每个节点i （无论是源令牌还是目标令牌），您都可以将注意力计算分解为两个步骤：
- 消息计算：计算注意力得分 scoreij 之间 i 和所有节点 j 通过参加之间的缩放点积来参加 qi 和 kj。邮件发自j 至 i 将由分数组成 scoreij 和价值 vj。
- 消息聚合：聚合值vj 来自所有 j 根据分数 scoreij。
Simple implementation
Message computation
计算score源节点的并将其发送v到目标邮箱
def message_func(edges):
return {'score': ((edges.src['k'] * edges.dst['q'])
.sum(-1, keepdim=True)),
'v': edges.src['v']}
Message aggregation
对所有边缘和加权和进行归一化以获得输出
import torch as th
import torch.nn.functional as F
def reduce_func(nodes, d_k=64):
v = nodes.mailbox['v']
att = F.softmax(nodes.mailbox['score'] / th.sqrt(d_k), 1)
return {'dx': (att * v).sum(1)}
Execute on specific edges
import functools.partial as partial
def naive_propagate_attention(self, g, eids):
g.send_and_recv(eids, message_func, partial(reduce_func, d_k=self.d_k))
Speeding up with built-in functions
为了加快消息传递过程，请使用DGL的内置功能，包括：
- fn.src_mul_egdes(src_field, edges_field, out_field)将源属性和edges属性相乘，然后将结果发送到以键为键的目标节点的邮箱out_field。
- fn.copy_edge(edges_field, out_field) 将edge的属性复制到目标节点的邮箱.
- fn.sum(edges_field, out_field) 汇总边缘的属性，然后将聚合发送到目标节点的邮箱。
在这里，您将这些内置函数组装为propagate_attention，这也是最终实现中的主要图形操作函数。要使其加速，请将softmax操作分为以下步骤。回想一下，每个负责人都有两个阶段。
1、通过将src节点k和dst节点的 乘积计算注意力分数q
- g.apply_edges(src_dot_dst(‘k’, ‘q’, ‘score’), eids)
2、所有dst节点的传入边缘上缩放的Softmax
第1步：使用规模归一化常数对指数进行指数化
-
g.apply_edges(scaled_exp(‘score’, np.sqrt(self.d_k)))
scoreij←exp(dkscoreij)
第2步：获取关联节点上的“值”，并按每个节点的传入边缘上的“得分”加权；获取每个节点的传入边缘上的“分数”之和以进行标准化。注意这里 wv 未标准化。
-
msg: fn.src_mul_edge(‘v’, ‘score’, ‘v’), reduce: fn.sum(‘v’, ‘wv’)
wvj=i=1∑Nscoreij⋅vi -
msg: fn.copy_edge(‘score’, ‘score’), reduce: fn.sum(‘score’, ‘z’)
zj=i=1∑Nscoreij
归一化 wv留给后期处理。
def src_dot_dst(src_field, dst_field, out_field):
def func(edges):
return {out_field: (edges.src[src_field] * edges.dst[dst_field]).sum(-1, keepdim=True)}
return func
def scaled_exp(field, scale_constant):
def func(edges):
# clamp for softmax numerical stability
return {field: th.exp((edges.data[field] / scale_constant).clamp(-5, 5))}
return func
def propagate_attention(self, g, eids):
# Compute attention score
g.apply_edges(src_dot_dst('k', 'q', 'score'), eids)
g.apply_edges(scaled_exp('score', np.sqrt(self.d_k)))
# Update node state
g.send_and_recv(eids,
[fn.src_mul_edge('v', 'score', 'v'), fn.copy_edge('score', 'score')],
[fn.sum('v', 'wv'), fn.sum('score', 'z')])
Preprocessing and postprocessing
在Transformer中，数据需要在propagate_attention函数之前和之后进行预处理。
预处理预处理功能pre_func首先规范化节点表示形式，然后以自我关注为例将它们映射到一组查询，键和值：
x←LayerNorm(x)[q,k,v]←[Wq,Wk,Wv]⋅x
后处理后处理功能post_funcs完成对应于变压器一层的整个计算：1.规范化wv 并获得多头注意力层的输出 o。
wv←zwvo←Wo⋅wv+bo
添加剩余连接：
x←x+o
2、在其上应用两层位置前馈层 x 然后添加剩余连接：
x←x+LayerNorm(FFN(x))
哪里 FFN 指前馈功能。
class Encoder(nn.Module):
def __init__(self, layer, N):
super(Encoder, self).__init__()
self.N = N
self.layers = clones(layer, N)
self.norm = LayerNorm(layer.size)
def pre_func(self, i, fields='qkv'):
layer = self.layers[i]
def func(nodes):
x = nodes.data['x']
norm_x = layer.sublayer[0].norm(x)
return layer.self_attn.get(norm_x, fields=fields)
return func
def post_func(self, i):
layer = self.layers[i]
def func(nodes):
x, wv, z = nodes.data['x'], nodes.data['wv'], nodes.data['z']
o = layer.self_attn.get_o(wv / z)
x = x + layer.sublayer[0].dropout(o)
x = layer.sublayer[1](x, layer.feed_forward)
return {'x': x if i < self.N - 1 else self.norm(x)}
return func
class Decoder(nn.Module):
def __init__(self, layer, N):
super(Decoder, self).__init__()
self.N = N
self.layers = clones(layer, N)
self.norm = LayerNorm(layer.size)
def pre_func(self, i, fields='qkv', l=0):
layer = self.layers[i]
def func(nodes):
x = nodes.data['x']
if fields == 'kv':
norm_x = x # In enc-dec attention, x has already been normalized.
else:
norm_x = layer.sublayer[l].norm(x)
return layer.self_attn.get(norm_x, fields)
return func
def post_func(self, i, l=0):
layer = self.layers[i]
def func(nodes):
x, wv, z = nodes.data['x'], nodes.data['wv'], nodes.data['z']
o = layer.self_attn.get_o(wv / z)
x = x + layer.sublayer[l].dropout(o)
if l == 1:
x = layer.sublayer[2](x, layer.feed_forward)
return {'x': x if i < self.N - 1 else self.norm(x)}
return func
这样就完成了Transformer中一层编码器和解码器的所有过程。
Note:
子层连接部分与原始纸张略有不同。但是，此实现与The Annotated Transformer 和OpenNMT相同。
Main class of Transformer graph
可以将Transformer的处理流程视为完整图形中的两阶段消息传递（适当地添加预处理和后处理）：1）编码器中的自注意力，2）解码器中的自注意力，然后交叉编码器和解码器之间的注意，如下所示。
图片地址：https://i.imgur.com/zlUpJ41.png
class Transformer(nn.Module):
def __init__(self, encoder, decoder, src_embed, tgt_embed, pos_enc, generator, h, d_k):
super(Transformer, self).__init__()
self.encoder, self.decoder = encoder, decoder
self.src_embed, self.tgt_embed = src_embed, tgt_embed
self.pos_enc = pos_enc
self.generator = generator
self.h, self.d_k = h, d_k
def propagate_attention(self, g, eids):
# Compute attention score
g.apply_edges(src_dot_dst('k', 'q', 'score'), eids)
g.apply_edges(scaled_exp('score', np.sqrt(self.d_k)))
# Send weighted values to target nodes
g.send_and_recv(eids,
[fn.src_mul_edge('v', 'score', 'v'), fn.copy_edge('score', 'score')],
[fn.sum('v', 'wv'), fn.sum('score', 'z')])
def update_graph(self, g, eids, pre_pairs, post_pairs):
"Update the node states and edge states of the graph."
# Pre-compute queries and key-value pairs.
for pre_func, nids in pre_pairs:
g.apply_nodes(pre_func, nids)
self.propagate_attention(g, eids)
# Further calculation after attention mechanism
for post_func, nids in post_pairs:
g.apply_nodes(post_func, nids)
def forward(self, graph):
g = graph.g
nids, eids = graph.nids, graph.eids
# Word Embedding and Position Embedding
src_embed, src_pos = self.src_embed(graph.src[0]), self.pos_enc(graph.src[1])
tgt_embed, tgt_pos = self.tgt_embed(graph.tgt[0]), self.pos_enc(graph.tgt[1])
g.nodes[nids['enc']].data['x'] = self.pos_enc.dropout(src_embed + src_pos)
g.nodes[nids['dec']].data['x'] = self.pos_enc.dropout(tgt_embed + tgt_pos)
for i in range(self.encoder.N):
# Step 1: Encoder Self-attention
pre_func = self.encoder.pre_func(i, 'qkv')
post_func = self.encoder.post_func(i)
nodes, edges = nids['enc'], eids['ee']
self.update_graph(g, edges, [(pre_func, nodes)], [(post_func, nodes)])
for i in range(self.decoder.N):
# Step 2: Dncoder Self-attention
pre_func = self.decoder.pre_func(i, 'qkv')
post_func = self.decoder.post_func(i)
nodes, edges = nids['dec'], eids['dd']
self.update_graph(g, edges, [(pre_func, nodes)], [(post_func, nodes)])
# Step 3: Encoder-Decoder attention
pre_q = self.decoder.pre_func(i, 'q', 1)
pre_kv = self.decoder.pre_func(i, 'kv', 1)
post_func = self.decoder.post_func(i, 1)
nodes_e, nodes_d, edges = nids['enc'], nids['dec'], eids['ed']
self.update_graph(g, edges, [(pre_q, nodes_d), (pre_kv, nodes_e)], [(post_func, nodes_d)])
return self.generator(g.ndata['x'][nids['dec']])
Note:
通过调用update_graphfunction，您可以使用几乎相同的代码在任何子图中创建自己的Transformer。这种灵活性使我们能够发现新的稀疏结构（请参见此处提到的本地关注）。请注意，在此实现中，您不使用掩码或填充，这使逻辑更加清晰并节省了内存。权衡是实施速度较慢。
Training
本教程没有涵盖其他技术，例如原始论文中提到的标签平滑和Noam优化。有关这些模块的详细说明，请阅读 由哈佛NLP团队编写的带注释的变压器。
Task and the dataset
变压器是用于各种NLP任务的通用框架。本教程重点介绍序列学习的序列：这是说明其工作原理的典型案例。
对于数据集，有两个示例任务：复制和排序，以及两个实际翻译任务：multi30k en-de任务和wmt14 en-de任务。
- 复制数据集：将输入序列复制到输出。（训练/有效/测试：9000、1000、1000）
- 排序数据集：将输入序列排序为输出。（训练/有效/测试：9000、1000、1000）
- Multi30k en-de，将句子从En转换为De。（训练/有效/测试：29000、1000、1000）
- WMT14 en-de，从En到De翻译句子。（训练/有效/测试：4500966/3000/3003）
Note:
使用wmt14进行培训需要多GPU支持，并且不可用。欢迎做贡献！
Graph building
批处理这类似于您处理Tree-LSTM的方式。预先构建一个图形池，包括输入长度和输出长度的所有可能组合。然后，对于批次中的每个样本，将dgl.batch其大小的批次图一起调用为一个大图。
您可以在dataset.GraphPool和中 包装创建图形池和构建BatchedGraph的过程dataset.TranslationDataset。
graph_pool = GraphPool()
data_iter = dataset(graph_pool, mode='train', batch_size=1, devices=devices)
for graph in data_iter:
print(graph.nids['enc']) # encoder node ids
print(graph.nids['dec']) # decoder node ids
print(graph.eids['ee']) # encoder-encoder edge ids
print(graph.eids['ed']) # encoder-decoder edge ids
print(graph.eids['dd']) # decoder-decoder edge ids
print(graph.src[0]) # Input word index list
print(graph.src[1]) # Input positions
print(graph.tgt[0]) # Output word index list
print(graph.tgt[1]) # Ouptut positions
break
out:
tensor([0, 1, 2, 3, 4, 5, 6, 7, 8], device='cuda:0')
tensor([ 9, 10, 11, 12, 13, 14, 15, 16, 17, 18], device='cuda:0')
tensor([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17,
18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,
36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53,
54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71,
72, 73, 74, 75, 76, 77, 78, 79, 80], device='cuda:0')
tensor([ 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94,
95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108,
109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122,
123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136,
137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150,
151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164,
165, 166, 167, 168, 169, 170], device='cuda:0')
tensor([171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184,
185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198,
199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210, 211, 212,
213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225],
device='cuda:0')
tensor([28, 25, 7, 26, 6, 4, 5, 9, 18], device='cuda:0')
tensor([0, 1, 2, 3, 4, 5, 6, 7, 8], device='cuda:0')
tensor([ 0, 28, 25, 7, 26, 6, 4, 5, 9, 18], device='cuda:0')
tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], device='cuda:0')
Put it all together
在复制任务上训练一层128尺寸的单头变压器。将其他参数设置为默认值。
推理模块未包含在本教程中。它需要光束搜索。有关完整的实现，请参见GitHub repo。
from tqdm import tqdm
import torch as th
import numpy as np
from loss import LabelSmoothing, SimpleLossCompute
from modules import make_model
from optims import NoamOpt
from dgl.contrib.transformer import get_dataset, GraphPool
def run_epoch(data_iter, model, loss_compute, is_train=True):
for i, g in tqdm(enumerate(data_iter)):
with th.set_grad_enabled(is_train):
output = model(g)
loss = loss_compute(output, g.tgt_y, g.n_tokens)
print('average loss: {}'.format(loss_compute.avg_loss))
print('accuracy: {}'.format(loss_compute.accuracy))
N = 1
batch_size = 128
devices = ['cuda' if th.cuda.is_available() else 'cpu']
dataset = get_dataset("copy")
V = dataset.vocab_size
criterion = LabelSmoothing(V, padding_idx=dataset.pad_id, smoothing=0.1)
dim_model = 128
# Create model
model = make_model(V, V, N=N, dim_model=128, dim_ff=128, h=1)
# Sharing weights between Encoder & Decoder
model.src_embed.lut.weight = model.tgt_embed.lut.weight
model.generator.proj.weight = model.tgt_embed.lut.weight
model, criterion = model.to(devices[0]), criterion.to(devices[0])
model_opt = NoamOpt(dim_model, 1, 400,
th.optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.98), eps=1e-9))
loss_compute = SimpleLossCompute
att_maps = []
for epoch in range(4):
train_iter = dataset(graph_pool, mode='train', batch_size=batch_size, devices=devices)
valid_iter = dataset(graph_pool, mode='valid', batch_size=batch_size, devices=devices)
print('Epoch: {} Training...'.format(epoch))
model.train(True)
run_epoch(train_iter, model,
loss_compute(criterion, model_opt), is_train=True)
print('Epoch: {} Evaluating...'.format(epoch))
model.att_weight_map = None
model.eval()
run_epoch(valid_iter, model,
loss_compute(criterion, None), is_train=False)
att_maps.append(model.att_weight_map)
Visualization
训练后，您可以形象地看到Transformer在复制任务上产生的注意力。
src_seq = dataset.get_seq_by_id(VIZ_IDX, mode='valid', field='src')
tgt_seq = dataset.get_seq_by_id(VIZ_IDX, mode='valid', field='tgt')[:-1]
# visualize head 0 of encoder-decoder attention
att_animation(att_maps, 'e2d', src_seq, tgt_seq, 0)
从图中可以看到，解码器节点逐渐学习按照输入顺序参与相应的节点，这是预期的行为。
Multi-head attention
除了专注于玩具任务的单头注意力训练。我们还可视化了在多30k数据集上训练的单层变压器网络的编码器的自我注意，解码器的自我注意和编码器-解码器注意的注意力得分。
从可视化中，您可以看到不同头的多样性。不同的头脑学习单词对之间的不同关系。
- 编码器自我注意
图片地址：https://i.imgur.com/HjYb7F2.png
- 编码器-解码器注意目标序列中的大多数单词都与源序列中它们的相关单词相关，例如：生成“ See”（在De中）时，多个头部都在“ lake”上；生成“Eisfischerhütte”时，几个负责人参加了“ ice”。
图片地址：https://i.imgur.com/383J5O5.png
- 解码器的自我注意大多数单词都以其前几个单词出现。
图片地址：https://i.imgur.com/c0UWB1V.png
Adaptive Universal Transformer
谷歌最近发表的研究论文Universal Transformer就是一个例子，展示了如何update_graph适应更复杂的更新规则。
提出通用变压器是为了通过在Transformer中引入递归来解决香草变压器在计算上不通用的问题：
- Universal Transformer的基本思想是，在每个重复步骤中，通过在表示形式上施加一个Transformer层，反复修改其在序列中所有符号的表示形式。
- 与普通Transformer相比，Universal Transformer在其各层之间共享权重，并且它不固定重复时间（这意味着Transformer中的层数）。
进一步的优化采用自适应计算时间（ACT）机制，以允许模型动态调整序列中每个位置的表示被修改的次数（ 此后称为步骤）。该模型也称为自适应通用变压器（AUT）。
在AUT中，您维护一个活动节点列表。在每一步t，我们计算出停止的概率： h(0<h<1) 对于此列表中的所有节点，通过：
hit=σ(Whxit+bh)
然后动态决定哪些节点仍处于活动状态。某个节点暂停T 当且仅当 ∑t=1T−1ht<1−ε≤∑t=1Tht。暂停的节点将从列表中删除。该过程将继续进行，直到列表为空或达到预定义的最大步骤为止。从DGL的角度来看，这意味着“活动”图随时间变得稀疏。
节点的最终状态 si 是的加权平均值 xti 通过 hti：
si=t=1∑Thit⋅xit
在DGL中，通过调用update_graph仍处于活动状态的节点以及与此节点关联的边来实现算法 。以下代码显示了DGL中的Universal Transformer类：
class UTransformer(nn.Module):
"Universal Transformer(https://arxiv.org/pdf/1807.03819.pdf) with ACT(https://arxiv.org/pdf/1603.08983.pdf)."
MAX_DEPTH = 8
thres = 0.99
act_loss_weight = 0.01
def __init__(self, encoder, decoder, src_embed, tgt_embed, pos_enc, time_enc, generator, h, d_k):
super(UTransformer, self).__init__()
self.encoder, self.decoder = encoder, decoder
self.src_embed, self.tgt_embed = src_embed, tgt_embed
self.pos_enc, self.time_enc = pos_enc, time_enc
self.halt_enc = HaltingUnit(h * d_k)
self.halt_dec = HaltingUnit(h * d_k)
self.generator = generator
self.h, self.d_k = h, d_k
def step_forward(self, nodes):
# add positional encoding and time encoding, increment step by one
x = nodes.data['x']
step = nodes.data['step']
pos = nodes.data['pos']
return {'x': self.pos_enc.dropout(x + self.pos_enc(pos.view(-1)) + self.time_enc(step.view(-1))),
'step': step + 1}
def halt_and_accum(self, name, end=False):
"field: 'enc' or 'dec'"
halt = self.halt_enc if name == 'enc' else self.halt_dec
thres = self.thres
def func(nodes):
p = halt(nodes.data['x'])
sum_p = nodes.data['sum_p'] + p
active = (sum_p < thres) & (1 - end)
_continue = active.float()
r = nodes.data['r'] * (1 - _continue) + (1 - sum_p) * _continue
s = nodes.data['s'] + ((1 - _continue) * r + _continue * p) * nodes.data['x']
return {'p': p, 'sum_p': sum_p, 'r': r, 's': s, 'active': active}
return func
def propagate_attention(self, g, eids):
# Compute attention score
g.apply_edges(src_dot_dst('k', 'q', 'score'), eids)
g.apply_edges(scaled_exp('score', np.sqrt(self.d_k)), eids)
# Send weighted values to target nodes
g.send_and_recv(eids,
[fn.src_mul_edge('v', 'score', 'v'), fn.copy_edge('score', 'score')],
[fn.sum('v', 'wv'), fn.sum('score', 'z')])
def update_graph(self, g, eids, pre_pairs, post_pairs):
"Update the node states and edge states of the graph."
# Pre-compute queries and key-value pairs.
for pre_func, nids in pre_pairs:
g.apply_nodes(pre_func, nids)
self.propagate_attention(g, eids)
# Further calculation after attention mechanism
for post_func, nids in post_pairs:
g.apply_nodes(post_func, nids)
def forward(self, graph):
g = graph.g
N, E = graph.n_nodes, graph.n_edges
nids, eids = graph.nids, graph.eids
# embed & pos
g.nodes[nids['enc']].data['x'] = self.src_embed(graph.src[0])
g.nodes[nids['dec']].data['x'] = self.tgt_embed(graph.tgt[0])
g.nodes[nids['enc']].data['pos'] = graph.src[1]
g.nodes[nids['dec']].data['pos'] = graph.tgt[1]
# init step
device = next(self.parameters()).device
g.ndata['s'] = th.zeros(N, self.h * self.d_k, dtype=th.float, device=device) # accumulated state
g.ndata['p'] = th.zeros(N, 1, dtype=th.float, device=device) # halting prob
g.ndata['r'] = th.ones(N, 1, dtype=th.float, device=device) # remainder
g.ndata['sum_p'] = th.zeros(N, 1, dtype=th.float, device=device) # sum of pondering values
g.ndata['step'] = th.zeros(N, 1, dtype=th.long, device=device) # step
g.ndata['active'] = th.ones(N, 1, dtype=th.uint8, device=device) # active
for step in range(self.MAX_DEPTH):
pre_func = self.encoder.pre_func('qkv')
post_func = self.encoder.post_func()
nodes = g.filter_nodes(lambda v: v.data['active'].view(-1), nids['enc'])
if len(nodes) == 0: break
edges = g.filter_edges(lambda e: e.dst['active'].view(-1), eids['ee'])
end = step == self.MAX_DEPTH - 1
self.update_graph(g, edges,
[(self.step_forward, nodes), (pre_func, nodes)],
[(post_func, nodes), (self.halt_and_accum('enc', end), nodes)])
g.nodes[nids['enc']].data['x'] = self.encoder.norm(g.nodes[nids['enc']].data['s'])
for step in range(self.MAX_DEPTH):
pre_func = self.decoder.pre_func('qkv')
post_func = self.decoder.post_func()
nodes = g.filter_nodes(lambda v: v.data['active'].view(-1), nids['dec'])
if len(nodes) == 0: break
edges = g.filter_edges(lambda e: e.dst['active'].view(-1), eids['dd'])
self.update_graph(g, edges,
[(self.step_forward, nodes), (pre_func, nodes)],
[(post_func, nodes)])
pre_q = self.decoder.pre_func('q', 1)
pre_kv = self.decoder.pre_func('kv', 1)
post_func = self.decoder.post_func(1)
nodes_e = nids['enc']
edges = g.filter_edges(lambda e: e.dst['active'].view(-1), eids['ed'])
end = step == self.MAX_DEPTH - 1
self.update_graph(g, edges,
[(pre_q, nodes), (pre_kv, nodes_e)],
[(post_func, nodes), (self.halt_and_accum('dec', end), nodes)])
g.nodes[nids['dec']].data['x'] = self.decoder.norm(g.nodes[nids['dec']].data['s'])
act_loss = th.mean(g.ndata['r']) # ACT loss
return self.generator(g.ndata['x'][nids['dec']]), act_loss * self.act_loss_weight
调用filter_nodes并filter_edge查找仍处于活动状态的节点/边：
Note:
- filter_nodes() 将谓词和节点ID列表/张量作为输入，然后返回满足给定谓词的节点ID的张量。
- filter_edges() 接受谓词和边缘ID列表/张量作为输入，然后返回满足给定谓词的边缘ID的张量。
有关完整的实现，请参见GitHub repo。
下图显示了自适应计算时间的影响。句子的不同位置在不同的时间进行了修订。
您还可以在排序任务的AUT训练过程中可视化节点上步长分布的动态（达到99.7％的准确性），这演示了AUT如何学习减少训练中的重复步骤。
Note:
由于存在许多依赖性，笔记本本身无法执行。下载7_transformer.py，并将python脚本复制到目录，examples/pytorch/transformer
然后运行以查看其工作方式。python 7_transformer.py
脚本的总运行时间：（0分钟0.063秒）
下载脚本：7_transformer.py
下载脚本：7_transformer.ipynb