I implemented a method fit to train the model that uses nn.Transformer with the teacher-forcing approach.
Unfortunately, I noticed that the loss was not getting small, and the inference forward was really bad.
I tried increasing the token's dimension size, but it didn't help.
Would you guys check my code to see if I misunderstood something?
I documented the code as much as possible for readability.
Here is my code (this is the initial code I wrote, and I wrote another one below and updated this post):
import torch
import torch.nn as nn
import torch.optim as optim
seed = 3
device = torch.device("cpu")
torch.manual_seed(seed)
if torch.cuda.is_available():
device = torch.device("cuda")
torch.cuda.manual_seed_all(seed)
class OurModel(nn.Module):
def __init__(self, token_to_index):
super().__init__()
self.token_to_index = token_to_index
self.index_to_token = {index: token for token, index in token_to_index.items()}
self.vocab_size = len(token_to_index)
# FYI, 512 is the default for the transformer.
token_vec_dim = 2
self.embeds = nn.Embedding(self.vocab_size, token_vec_dim)
self.transformer = nn.Transformer(
d_model=token_vec_dim,
# `d_model` must be divisible by `nhead`.
# FYI, default `nhead` is 8.
# Also, you will get warnings if you set `nhead` to an odd number.
nhead=8 if token_vec_dim % 8 == 0 else 2,
batch_first=True
)
self.fc = nn.Sequential(
nn.Linear(token_vec_dim, self.vocab_size)
)
def forward(self, sequence, max_iteration=10):
"""
Generate a sequence of tokens in an autoregressive way.
The generation ends when <eos> token is generated or reaches the maximum iteration.
Here is how it works:
First, we convert tokens into vectors.
Second, we need two parameters to use `nn.Transformer(src, tgt)`.
The `src` represents the input vectors.
Ex: A tensor (of vectors) for ["roses", "are"].
The `tgt` represents the generated vectors so far.
Initially, we use the tensor for ["<sos>"].
The output of `nn.Transformer(src, tgt)` is a tensor with the same shape as the `tgt`.
For example, if `tgt` was a tensor for ["<sos>"], the output will be a tensor of vectors like [[1, 2]].
We grab the [1, 2], put it into the fully connected layer to predict the next token.
Once we got the next token, we append it to the `tgt`.
For example, the `tgt` for the next iteration could be a tensor for ["<sos>", "red"].
The transformer's output may look like a tensor of vectors like [[1, 2], [3, 4]].
Note that we only grab the last element [3, 4] and ignore the rest.
We put the last element into the fully connected layer to predict the next token.
If the predicted token is "<eos>", we stop the generation.
:param sequence: List of tokens. Ex: ["roses", "are"].
:param max_iteration: We stop generating tokens if it reaches the max_interation without getting the "<eos>" token.
:return: List of tokens. Ex: ["<sos>", "red", "<eos>"].
"""
token_indexes_raw = [self.token_to_index[token] for token in sequence]
# Shape: (token_amount).
token_indexes_tensor = torch.LongTensor([index for index in token_indexes_raw]).to(device)
# Shape: (token_amount, token_vec_dim).
token_vectors = self.embeds(token_indexes_tensor)
# Shape: (1, token_vec_dim).
sos_token_vec = self.embeds(torch.LongTensor([0]).to(device))
# Shape: (generated_token_amount_so_far, token_vec_dim).
current_target = sos_token_vec
generated_token_indexes = [0]
iteration = 1
predicted_token_index = -1
while predicted_token_index != 1 and iteration < max_iteration:
# The shape is the same as the target.
transformer_output = self.transformer(token_vectors, current_target)
# Shape: (token_vec_dim).
last_element = transformer_output[-1]
# Shape: (vocab_size)
raw_scores = self.fc(last_element)
predicted_token_index = torch.argmax(raw_scores).item()
generated_token_indexes.append(predicted_token_index)
# Shape: (1, token_vec_dim).
last_generated_token_vec = self.embeds(torch.LongTensor([predicted_token_index]).to(device))
current_target = torch.cat((current_target, last_generated_token_vec), dim=0).to(device)
iteration += 1
return [self.index_to_token[index] for index in generated_token_indexes]
def fit(self, sources, targets):
"""
Train the model.
Here is how it works:
First, we convert tokens into vectors.
Second, we get the output from `nn.Transformer(sources, targets)`.
The transformer's output will be a batch of generated tokens.
Note that the output's shape is the same as the shape of targets.
Then, we put the transformer's output as an input for the fully connected layer.
The fully connected layer will give a batch of raw scores.
With each raw scores, we can determine the predicted next token.
For the loss, we can compare the predicted next token and the actual next token from the target of the transformer.
For example, the targets were: [["<sos>", "red"], ["<sos>", "green"]].
The transformer's output was: [[value_1, value_2], [value_3, value_4]].
The fully connected layer's output was: [[raw_score_1, raw_score_2], [raw_score_3, raw_score_4]].
The raw_score_1 should've predicted the token "red".
The raw_score_2 should've predicted the token "<eos>".
As you can see, the last score corresponds to the "<eos>" token.
That's why we excluded the "<eos>" token in targets.
:param sources: A batch of sources in a plain list. Ex: [["roses", "are"], ["limes", "are"]].
:param targets: A batch of targets in a plain list. Ex: [["red"], ["green"]].
Actually, it's supposed to look like this: [["<sos>", "red"], ["<sos>", "green"]].
However, this method will put "<sos>" for you.
Note that we exclude "<eos>" on purpose.
:return: None.
"""
loss_function = nn.CrossEntropyLoss()
optimizer = optim.Adam(self.parameters())
for epoch in range(300):
batch_source_indexes = [[self.token_to_index[token] for token in source] for source in sources]
# Shape: (batch_size, token_amount_in_sequence, token_vec_dim).
batch_source_vectors = self.embeds(torch.LongTensor(batch_source_indexes).to(device))
# Note that we add the index of the "<sos>" token at the beginning for each target token.
batch_target_indexes = [[0] + [self.token_to_index[token] for token in target] for target in targets]
# Shape: (batch_size, token_amount_in_sequence, token_vec_dim).
batch_target_vectors = self.embeds(torch.LongTensor(batch_target_indexes).to(device))
# The shape is the same as the target.
transformer_output = self.transformer(batch_source_vectors, batch_target_vectors)
# Shape: (batch_size, token_amount_in_sequence, vocab_size).
batch_raw_scores = self.fc(transformer_output)
# Note that we add the index of the "<eos>" token at the end for each target token.
batch_correct_token_indexes = [[self.token_to_index[token] for token in target] + [1] for target in targets]
# Reshape to (batch_size * token_amount_in_sequence, vocab_size).
# The reason is that the loss function requires 2D tensor as an input.
# It's okay to reshape it because it's effectively the same thing as iterating the batch.
# view method is the same as reshape method, except it's mutable.
# The first -1 is a magic number for automatically calculate the appropriate size of the first dimension based on other dimension.
# The second -1 means the size of the last dimension.
# Let's say the original shape is (4, 2, 9).
# The second -1 determined the size of the last dimension, which is 9.
# That makes the remaining dimension to be 4 * 2 = 8 because view and reshape methods keep the number of data.
# As a result, the new shape will be (8, 9).
flatten_raw_scores = batch_raw_scores.view(-1, batch_raw_scores.size(-1))
# Shape: (batch_size * token_amount_in_sequence).
flatten_correct_token_indexes_tensor = torch.LongTensor(batch_correct_token_indexes).to(device).view(-1)
loss = loss_function(flatten_raw_scores, flatten_correct_token_indexes_tensor)
print(f"epoch: {epoch}, loss: {loss}")
loss.backward()
optimizer.step()
optimizer.zero_grad()
def main():
token_to_index = {
# sos stands for start of sequence.
"<sos>": 0,
# eos stands for end of sequence.
"<eos>": 1,
"roses": 2,
"apples": 3,
"limes": 4,
"cucumbers": 5,
"are": 6,
"red": 7,
"green": 8
}
model = OurModel(token_to_index).to(device)
train_sources = [
["roses", "are"],
["apples", "are"],
["limes", "are"],
["cucumbers", "are"]
]
train_targets = [
["red"],
["red"],
["green"],
["green"]
]
model.fit(train_sources, train_targets)
input_tokens = ["roses", "are"]
generated_tokens = model(input_tokens)
print(f"input_tokens: {input_tokens}, generated_tokens: {generated_tokens}")
main()
EDIT:
I rewrote the code from scratch again just for retrying purposes, and I noticed that the loss was steadily decreasing. Maybe there were some mistakes in my first code.
Here is my second code (it has a training and inference part in the main function):
import torch
import torch.nn as nn
import torch.optim as optim
seed = 3
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
class OurModel(nn.Module):
def __init__(self, token_to_index):
super().__init__()
self.token_to_index = token_to_index
self.index_to_token = {index: token for token, index in token_to_index.items()}
self.vocab_size = len(token_to_index)
token_vec_dim = 16
self.embeds = nn.Embedding(self.vocab_size, token_vec_dim)
self.transformer = nn.Transformer(
d_model=token_vec_dim,
nhead=8 if token_vec_dim % 8 == 0 else 2,
batch_first=True
)
self.fc = nn.Sequential(
nn.Linear(token_vec_dim, 30),
nn.ReLU(),
nn.Linear(30, self.vocab_size)
)
def main():
token_to_index = {
"<sos>": 0,
"<eos>": 1,
"roses": 2,
"are": 3,
"red": 4,
"cats": 5,
"eat": 6,
"tuna": 7
}
model = OurModel(token_to_index)
loss_function = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())
for epoch in range(500):
# Put indexes for ["roses", "are"].
# Shape: (number_of_input_tokens, token_vec_dim).
input_token_vecs = model.embeds(torch.LongTensor([2, 3]))
# Represents ["<sos>", "red", "<eos>"].
# Shape: (number_of_target_tokens, token_vec_dim).
output_token_vecs = model.embeds(torch.LongTensor([0, 4, 1]))
# Shape: (number_of_target_tokens, token_vec_dim).
transformer_output = model.transformer(input_token_vecs, output_token_vecs)
# Shape: (number_of_target_tokens, vocab_size).
fc_output = model.fc(transformer_output)
# Since we don't have to predict anything after the <"eos"> token, we ignore the last element.
trimmed_fc_output = fc_output[:-1]
correct_indexes = torch.LongTensor([4, 1])
loss = loss_function(trimmed_fc_output, correct_indexes)
print(f"epoch: {epoch}, loss: {loss}")
if loss < 0.001:
break
loss.backward()
optimizer.step()
optimizer.zero_grad()
input_tokens = ["roses", "are"]
# Shape: (number_of_input_tokens).
input_token_indexes = torch.LongTensor([token_to_index[token] for token in input_tokens])
# Shape: (number_of_input_tokens, token_vec_dim).
input_token_vecs = model.embeds(input_token_indexes)
# Start with "<sos>" token for the target.
# Shape: (number_of_target_tokens).
target_token_indexes = torch.LongTensor([0])
next_token_index = -1
max_iteration = 10
iteration = 0
while next_token_index != 1 and iteration < max_iteration:
# Shape: (number_of_tokens, token_vec_dim).
target_token_vecs = model.embeds(target_token_indexes)
# Shape: (number_of_tokens, token_vec_dim).
transformer_output = model.transformer(input_token_vecs, target_token_vecs)
# Shape: (token_vec_dim).
last_element_of_transformer_output = transformer_output[-1]
# Shape: (vocab_size).
fc_output = model.fc(last_element_of_transformer_output)
# Shape: (0).
fc_output_argmax = torch.argmax(fc_output, dim=-1)
next_token_index = fc_output_argmax.item()
# Shape: (number_of_target_tokens).
target_token_indexes = torch.cat((target_token_indexes, fc_output_argmax.reshape(1)), dim=-1)
iteration += 1
print(f"Predicted output token indexes: {target_token_indexes}")
main()
token_vec_dim. $\endgroup$