Transformer多头注意例子

#首先，需要导入PyTorch和相关的模块：

import torch

import torch.nn as nn

import torch.optim as optim

import torch.nn.functional as F

#接着，定义Transformer模型的编码器（Encoder）和解码器（Decoder）：

class Encoder(nn.Module):

    def __init__(self, input_size, hidden_size, num_layers, num_heads, dropout):

        super(Encoder, self).__init__()

        self.input_size = input_size

        self.hidden_size = hidden_size

        self.num_layers = num_layers

        self.num_heads = num_heads

        self.dropout = dropout

        self.embedding = nn.Embedding(input_size, hidden_size)

        self.positional_encoding = PositionalEncoding(hidden_size, dropout)

        self.encoders = nn.ModuleList([EncoderLayer(hidden_size, num_heads, dropout) for _ in range(num_layers)])

    def forward(self, x):

        x = self.embedding(x) * math.sqrt(self.hidden_size)

        x = self.positional_encoding(x)

        for layer in self.encoders:

            x = layer(x)

        return x

class Decoder(nn.Module):

    def __init__(self, output_size, hidden_size, num_layers, num_heads, dropout):

        super(Decoder, self).__init__()

        self.output_size = output_size

        self.hidden_size = hidden_size

        self.num_layers = num_layers

        self.num_heads = num_heads

        self.dropout = dropout

        self.embedding = nn.Embedding(output_size, hidden_size)

        self.positional_encoding = PositionalEncoding(hidden_size, dropout)

        self.decoders = nn.ModuleList([DecoderLayer(hidden_size, num_heads, dropout) for _ in range(num_layers)])

    def forward(self, x, encoder_output):

        x = self.embedding(x) * math.sqrt(self.hidden_size)

        x = self.positional_encoding(x)

        for layer in self.decoders:

            x = layer(x, encoder_output)

        return x

#在上述代码中，我们使用了PyTorch中的Embedding层和ModuleList容器，以及自定义的位置编码（PositionalEncoding）、编码器层（EncoderLayer）和解码器层（DecoderLayer）。

#接着，定义位置编码层：

class PositionalEncoding(nn.Module):

    def __init__(self, hidden_size, dropout, max_length=5000):

        super(PositionalEncoding, self).__init__()

        self.hidden_size = hidden_size

        self.dropout = nn.Dropout(p=dropout)

        # 计算位置编码矩阵

        pe = torch.zeros(max_length, hidden_size)

        position = torch.arange(0, max_length, dtype=torch.float32).unsqueeze(1)

        div = torch.exp(torch.arange(0, hidden_size, 2).float() * (-math.log(10000.0) / hidden_size))

        pe[:, 0::2] = torch.sin(position * div)

        pe[:, 1::2] = torch.cos(position * div)

        pe = pe.unsqueeze(0).transpose(0, 1)

        self.register_buffer('pe', pe)

    def forward(self, x):

        x = x + self.pe[:x.size(0), :]

        return self.dropout(x)

#这里使用了PyTorch中的Dropout层。

#接下来，定义编码器层（EncoderLayer）和解码器层（DecoderLayer）：

class EncoderLayer(nn.Module):

    def __init__(self, hidden_size, num_heads, dropout):

        super(EncoderLayer, self).__init__()

        self.hidden_size = hidden_size

        self.num_heads = num_heads

        self.dropout = dropout

        self.self_attention = MultiheadAttention(hidden_size, num_heads, dropout)

        self.feedforward = FeedForward(hidden_size, dropout)

        self.layer_norm1 = nn.LayerNorm(hidden_size)

        self.layer_norm2 = nn.LayerNorm(hidden_size)

    def forward(self, x):

        x = self.layer_norm1(x + self.self_attention(x))

        x = self.layer_norm2(x + self.feedforward(x))

        return x

class DecoderLayer(nn.Module):

    def __init__(self, hidden_size, num_heads, dropout):

        super(DecoderLayer, self).__init__()

        self.hidden_size = hidden_size

        self.num_heads = num_heads

        self.dropout = dropout

        self.self_attention = MultiheadAttention(hidden_size, num_heads, dropout)

        self.encoder_attention = MultiheadAttention(hidden_size, num_heads, dropout)

        self.feedforward = FeedForward(hidden_size, dropout)

        self.layer_norm1 = nn.LayerNorm(hidden_size)

        self.layer_norm2 = nn.LayerNorm(hidden_size)

        self.layer_norm3 = nn.LayerNorm(hidden_size)

    def forward(self, x, encoder_output):

        x = self.layer_norm1(x + self.self_attention(x))

        x = self.layer_norm2(x + self.encoder_attention(x, encoder_output))

        x = self.layer_norm3(x + self.feedforward(x))

        return x

#在上述代码中，我们使用了自定义的多头注意力（MultiheadAttention）和前馈网络（FeedForward）。

#下面，定义多头注意力层和前馈网络层：

class MultiheadAttention(nn.Module):

    def __init__(self, hidden_size, num_heads, dropout):

        super(MultiheadAttention, self).__init__()

        self.hidden_size = hidden_size

        self.num_heads = num_heads

        self.dropout = dropout

        self.head_size = hidden_size // num_heads

        self.query = nn.Linear(hidden_size, hidden_size)

        self.key = nn.Linear(hidden_size, hidden_size)

        self.value = nn.Linear(hidden_size, hidden_size)

        self.fc = nn.Linear(hidden_size, hidden_size)

        self.dropout = nn.Dropout(p=dropout)

    def forward(self, x, enc_output=None):

        batch_size = x.size(0)

        q = self.query(x).view(batch_size, -1, self.num_heads, self.head_size).transpose(1, 2)

        k = self.key(x if enc_output is None else enc_output).view(batch_size, -1, self.num_heads, self.head_size).transpose(1, 2)

        v = self.value(x if enc_output is None else enc_output).view(batch_size, -1, self.num_heads, self.head_size).transpose(1, 2)

        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_size)

        scores = F.softmax(scores, dim=-1)

        scores = self.dropout(scores)

        x = torch.matmul(scores, v)

        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.hidden_size)

        x = self.fc(x)

        x = self.dropout(x)

        return x

class FeedForward(nn.Module):

    def __init__(self, hidden_size, dropout):

        super(FeedForward, self).__init__()

        self.hidden_size = hidden_size

        self.dropout = dropout

        self.fc1 = nn.Linear(hidden_size, hidden_size * 4)

        self.fc2 = nn.Linear(hidden_size * 4, hidden_size)

        self.dropout = nn.Dropout(p=dropout)

    def forward(self, x):

        x = self.fc1(x)

        x = F.relu(x)

        x = self.dropout(x)

        x = self.fc2(x)

        x = self.dropout(x)

        return x

#最后，定义完整的Transformer模型：

class Transformer(nn.Module):

    def __init__(self, input_size, output_size, hidden_size, num_layers, num_heads, dropout):

        super(Transformer, self).__init__()

        self.encoder = Encoder(input_size, hidden_size, num_layers, num_heads, dropout)

        self.decoder = Decoder(output_size, hidden_size, num_layers, num_heads, dropout)

        self.output_size = output_size

        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, input_seq, target_seq):

        encoder_output = self.encoder(input_seq)

        decoder_output = self.decoder(target_seq, encoder_output)

        decoder_output = self.fc(decoder_output)

        return decoder_output

#这里我们只定义了基本的Transformer模型，您可以根据具体的游戏场景进行个性化的修改和拓展，例如增加注意力权重控制等功能。同时，您也需要根据具体的游戏数据来定义数据预处理和模型训练的方式。

import random

import math

# 创建一些随机数据

def create_data(num_examples):

    input_seqs = []

    output_seqs = []

    for _ in range(num_examples):

        input_seq = [random.choice(range(1, 10)) for _ in range(5)]

        output_seq = sorted(input_seq)

        input_seqs.append(torch.tensor(input_seq, dtype=torch.long))

        output_seqs.append(torch.tensor(output_seq, dtype=torch.long))

    return input_seqs, output_seqs

# 定义训练函数

def train(model, data, epochs, device):

    model.train()

    criterion = nn.CrossEntropyLoss()

    optimizer = optim.Adam(model.parameters())

    for epoch in range(epochs):

        losses = []

        for input_seq, target_seq in data:

            input_seq = input_seq.to(device).unsqueeze(0)  # 增加一个维度

            target_seq = target_seq.to(device).unsqueeze(0)  # 增加一个维度

            optimizer.zero_grad()

            output = model(input_seq, target_seq[:, :-1])

            loss = criterion(output.view(-1, output.size(2)), target_seq[:, 1:].view(-1))

            loss.backward()

            optimizer.step()

            losses.append(loss.item())

        print(f'Epoch {epoch+1}/{epochs}, Loss: {sum(losses) / len(losses)}')

# 定义推理函数

def predict(model, input_seq, device):

    model.eval()

    input_seq = input_seq.to(device)

    output_seq = [0]

    for _ in range(5):

        target_seq = torch.tensor(output_seq, dtype=torch.long).unsqueeze(0).to(device)

        output = model(input_seq.unsqueeze(0), target_seq)

        output = output.argmax(2)[:, -1].item()

        if output != 0:

            output_seq.append(output)

    return output_seq[1:]

# 创建数据并将其分为训练和测试集

input_seqs, output_seqs = create_data(100)

train_data = list(zip(input_seqs[:-10], output_seqs[:-10]))

test_data = list(zip(input_seqs[-10:], output_seqs[-10:]))

# 初始化模型并将其移动到GPU（如有）

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = Transformer(input_size=11, output_size=11, hidden_size=512, num_layers=6, num_heads=8, dropout=0.1).to(device)

# 训练模型

train(model, train_data, epochs=20, device=device)

# 测试模型

for input_seq, target_seq in test_data[:5]:

    prediction = predict(model, input_seq, device)

    print(f'Input: {input_seq.tolist()}')

    print(f'True Output: {target_seq.tolist()}')

    print(f'Predicted Output: {prediction}')

    print()

#注意：上面的示例代码使用了一些随机生成的有序数据作为示例，实际应用中应该将其替换为特定游戏的数据。