gpt-4-0314· OpenAI WildChat-4.8M (AI2)

用pytorch搭建Transformer，完成一个简单任务，写出python代码

User
用pytorch搭建Transformer，完成一个简单任务，写出python代码

Assistant2023-04-09 06:40 UTC

这是一个使用PyTorch搭建Transformer并且完成一个序列到序列（seq2seq）任务的例子。我们将搭建一个简单的字符级别的神经网络翻译器，将输入的字符串翻译成另一种字符串形式。

import torch
import torch.nn as nn
import torch.optim as optim

from torchtext.data import Field, BucketIterator
from torchtext.datasets import TranslationDataset

import random
import math

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

# 将输入字符串翻译成另一种字符串形式的简单例子
# 输入字符串: "123"， 输出字符串: "one two three"
SRC = Field(tokenize=lambda x: list(x),
            init_token='<sos>',
            eos_token='<eos>',
            lower=True)
TRG = Field(tokenize=lambda x: x.split(),
            init_token='<sos>',
            eos_token='<eos>',
            lower=True)

raw_data = [('123', 'one two three'),
            ('456', 'four five six'),
            ('789', 'seven eight nine')]

examples = [TranslationDataset.Example.fromtuple(example, fields=[("src", SRC), ("trg", TRG)]) for example in raw_data]
dictionary = TranslationDataset(examples, fields=[("src", SRC), ("trg", TRG)])

# 构造用于训练的迭代器
SRC.build_vocab(dictionary, min_freq=1)
TRG.build_vocab(dictionary, min_freq=1)

BATCH_SIZE = 1

# 数据划分
train_data, valid_data, test_data = dictionary.split(split_ratio=[0.7,0.2,0.1])

train_iterator, valid_iterator, test_iterator = BucketIterator.splits(
    (train_data, valid_data, test_data),
    batch_size=BATCH_SIZE,
    device=device)

# Transformer 构造
class Transformer(nn.Module):
    def __init__(self, src_vocab_size, trg_vocab_size, d_model, nhead, num_layers):
        super().__init__()
        self.encoder_layer = nn.TransformerEncoderLayer(d_model, nhead)
        self.encoder = nn.TransformerEncoder(self.encoder_layer, num_layers)
        self.decoder_layer = nn.TransformerDecoderLayer(d_model, nhead)
        self.decoder = nn.TransformerDecoder(self.decoder_layer, num_layers)
        self.src_embedding = nn.Embedding(src_vocab_size, d_model)
        self.trg_embedding = nn.Embedding(trg_vocab_size, d_model)
        self.fc = nn.Linear(d_model, trg_vocab_size)

    def forward(self, src, trg):
        src_embedded = self.src_embedding(src) * math.sqrt(src.shape[1])
        trg_embedded = self.trg_embedding(trg) * math.sqrt(trg.shape[1])
        src_encoded = self.encoder(src_embedded)
        trg_decoded = self.decoder(trg_embedded, src_encoded)
        output = self.fc(trg_decoded)
        return output

# 参数
INPUT_DIM = len(SRC.vocab)
OUTPUT_DIM = len(TRG.vocab)
D_MODEL = 32
NHEAD = 2
NUM_LAYERS = 2

model = Transformer(INPUT_DIM, OUTPUT_DIM, D_MODEL, NHEAD, NUM_LAYERS).to(device)

# 训练函数
def train(model, iterator, optimizer, criterion, clip):
    model.train()
    epoch_loss = 0

    for batch in iterator:
        src = batch.src
        trg = batch.trg

        optimizer.zero_grad()

        output = model(src, trg[:-1, :])

        output_dim = output.shape[-1]

        output = output.view(-1, output_dim)
        trg = trg[1:].view(-1)

        loss = criterion(output, trg)
        loss.backward()

        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)

        optimizer.step()

        epoch_loss += loss.item()

    return epoch_loss / len(iterator)

def evaluate(model, iterator, criterion):
    model.eval()
    epoch_loss = 0

    with torch.no_grad():
        for batch in iterator:
            src = batch.src
            trg = batch.trg

            output = model(src, trg[:-1, :])

            output_dim = output.shape[-1]

            output = output.view(-1, output_dim)
            trg = trg[1:].view(-1)

            loss = criterion(output, trg)
            epoch_loss += loss.item()

    return epoch_loss / len(iterator)

N_EPOCHS = 5
CLIP = 1

optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss(ignore_index=TRG.vocab.stoi[TRG.pad_token])

best_valid_loss = float('inf')

for epoch in range(N_EPOCHS):
    train_loss = train(model, train_iterator, optimizer, criterion, CLIP)
    valid_loss = evaluate(model, valid_iterator, criterion)

    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        torch.save(model.state_dict(), 'model.pt')

    print(f'Epoch: {epoch+1:02}')
    print(f'\tTrain Loss: {train_loss:.3f}')
    print(f'\t Val. Loss: {valid_loss:.3f}')

model.load_state_dict(torch.load('model.pt'))
test_loss = evaluate(model, test_iterator, criterion)
print(f'| Test Loss: {test_loss:.3f} | Test PPL: {math.exp(test_loss):7.3f} |')

这个代码在5个epoch内训练一个简单的字符级别的神经网络翻译器。

User
以上代码报错：ImportError: cannot import name 'traverse_dps' from 'torch.utils.data.graph' ，请修改

Assistant2023-04-09 07:06 UTC

对不起给您带来不便。在此重新提供经过修改的代码，删除了traverse_dps相关的内容，并进行了相应的调整。

import torch
import torch.nn as nn
import torch.optim as optim
import random
import math

from torchtext.legacy.data import Field, BucketIterator, Example, Dataset

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

SRC = Field(tokenize=lambda x: list(x),
            init_token='<sos>',
            eos_token='<eos>',
            lower=True)

TRG = Field(tokenize=lambda x: x.split(),
            init_token='<sos>',
            eos_token='<eos>',
            lower=True)

raw_data = [('123', 'one two three'),
            ('456', 'four five six'),
            ('789', 'seven eight nine')]

examples = [Example.fromlist(example, fields=[("src", SRC), ("trg", TRG)]) for example in raw_data]
dictionary = Dataset(examples, fields=[("src", SRC), ("trg", TRG)])

SRC.build_vocab(dictionary, min_freq=1)
TRG.build_vocab(dictionary, min_freq=1)

BATCH_SIZE = 1
train_data, valid_data, test_data = dictionary.split(split_ratio=[0.7, 0.2, 0.1])

train_iterator, valid_iterator, test_iterator = BucketIterator.splits(
    (train_data, valid_data, test_data),
    batch_size=BATCH_SIZE,
    device=device)

class Transformer(nn.Module):
    def __init__(self, src_vocab_size, trg_vocab_size, d_model, nhead, num_layers):
        super().__init__()
        self.encoder_layer = nn.TransformerEncoderLayer(d_model, nhead)
        self.encoder = nn.TransformerEncoder(self.encoder_layer, num_layers)
        self.decoder_layer = nn.TransformerDecoderLayer(d_model, nhead)
        self.decoder = nn.TransformerDecoder(self.decoder_layer, num_layers)
        self.src_embedding = nn.Embedding(src_vocab_size, d_model)
        self.trg_embedding = nn.Embedding(trg_vocab_size, d_model)
        self.fc = nn.Linear(d_model, trg_vocab_size)

    def forward(self, src, trg):
        src_embedded = self.src_embedding(src) * math.sqrt(src.shape[1])
        trg_embedded = self.trg_embedding(trg) * math.sqrt(trg.shape[1])
        src_encoded = self.encoder(src_embedded)
        trg_decoded = self.decoder(trg_embedded, src_encoded)
        output = self.fc(trg_decoded)
        return output

INPUT_DIM = len(SRC.vocab)
OUTPUT_DIM = len(TRG.vocab)
D_MODEL = 32
NHEAD = 2
NUM_LAYERS = 2

model = Transformer(INPUT_DIM, OUTPUT_DIM, D_MODEL, NHEAD, NUM_LAYERS).to(device)

def train(model, iterator, optimizer, criterion, clip):
    model.train()
    epoch_loss = 0

    for batch in iterator:
        src = batch.src
        trg = batch.trg

        optimizer.zero_grad()

        output = model(src, trg[:-1, :])

        output_dim = output.shape[-1]

        output = output.view(-1, output_dim)
        trg = trg[1:].view(-1)

        loss = criterion(output, trg)
        loss.backward()

        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)

        optimizer.step()

        epoch_loss += loss.item()

    return epoch_loss / len(iterator)

def evaluate(model, iterator, criterion):
    model.eval()
    epoch_loss = 0

    with torch.no_grad():
        for batch in iterator:
            src = batch.src
            trg = batch.trg

            output = model(src, trg[:-1, :])

            output_dim = output.shape[-1]

            output = output.view(-1, output_dim)
            trg = trg[1:].view(-1)

            loss = criterion(output, trg)
            epoch_loss += loss.item()

    return epoch_loss / len(iterator)

N_EPOCHS = 5
CLIP = 1

optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss(ignore_index=TRG.vocab.stoi[TRG.pad_token])

best_valid_loss = float('inf')

for epoch in range(N_EPOCHS):
    train_loss = train(model, train_iterator, optimizer, criterion, CLIP)
    valid_loss = evaluate(model, valid_iterator, criterion)

    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        torch.save(model.state_dict(), 'model.pt')

    print(f'Epoch: {epoch+1:02}')
    print(f'\tTrain Loss: {train_loss:.3f}')
    print(f'\t Val. Loss: {valid_loss:.3f}')

model.load_state_dict(torch.load('model.pt'))
test_loss = evaluate(model, test_iterator, criterion)
print(f'| Test Loss: {test_loss:.3f} | Test PPL: {math.exp(test_loss):7.3f} |')

从torchtext 0.xx版本升级到torchtext 1.0(legacy)之后，上面的代码应该可以正常运行。如有任何问题，请继续留言。