如何让torch.nn.Transformer获得稳定的输出

Question

看起来 pytorch 的 Transformer 层无法提供可重现的输出。cpu和gpu都会出现这种情况。我知道有时会因为 GPU 上的并行计算而发生这种情况。

emb = nn.Embedding(10, 12).to(device)
inp1 = torch.LongTensor([1, 2, 3, 4]).to(device)
inp1 = emb(inp1).reshape(inp1.shape[0], 1, 12) #S N E

encoder_layer = nn.TransformerEncoderLayer(d_model=12, nhead=4)
transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=4)

out1 = transformer_encoder(inp1)
out2 = transformer_encoder(inp1)

out1 和 out2 不同。它可以在 cpu 上进行多处理，但结果看起来太不稳定。如何解决这个问题？

Answer 1

nn.TransformerEncoderLayer默认辍学率为0.1. 当模型处于训练模式时，要删除的索引将在每次迭代中随机化。

如果你想用 dropout 来训练模型，只需在训练中忽略这种行为并调用model.eval()测试即可。

如果你想在训练中禁用这种随机行为，请dropout=0像这样设置

nn.TransformerEncoderLayer(d_model=12, nhead=4, dropout=0)

完整的测试脚本：

import torch
import torch.nn as nn

device = 'cpu'

emb = nn.Embedding(10, 12).to(device)
inp1 = torch.LongTensor([1, 2, 3, 4]).to(device)
inp1 = emb(inp1).reshape(inp1.shape[0], 1, 12) #S N E

encoder_layer = nn.TransformerEncoderLayer(d_model=12, nhead=4, dropout=0).to(device)
transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=4).to(device)

out1 = transformer_encoder(inp1)
out2 = transformer_encoder(inp1)

print((out1-out2).norm())