I’ve been studying neural Transformer architecture for several months. Yesterday, I reached a major milestone when I successfully got a rudimentary prediction model running for the IMDB dataset to predict if a movie review is positive or negative.

I’ve spent several hundred hours on this knowledge quest, and I know I’ve got several hundred hours more ahead. I was ridiculously excited to get the crude prediction program running. The process made me think about the field of computer science, and learning in general.

The key to expertise in almost everything is relentless determination to succeed. This applies to music, athletics, writing, and so on, and especially to sciences like mathematics and computer science. In order to spend many thousands of hours studying something, I think a person has to be motivated internally.

This means that many efforts to push young people into computer science are doomed. If you convince a high school student or a college freshman that they should study computer science because the field pays well, you aren’t doing them any favors. It’s good to encourage students to explore many areas of study, then allow students to find areas that they’re passionate about. This will lead to long term success.

I’m often amused at efforts to push this group or that group in computer science. A direct analogy would be a giant effort to push men into the field of kindergarten teaching because they’re under represented there. Doesn’t make sense.

I’m fortunate that I wake up every morning, excited to get to work — and I mean that literally. I enjoy what I do because I tried many thing when I was young and found the things that I’m passionate about. If I worked in some other field, such as sales or retail, I doubt that I’d wake up as eager to get to work every morning as I am now.

My Transformer demo is very rough. It’s probably the single most complex program I’ve ever written. Even though it’s only a few hundred lines of code, almost every block of five or so statements is extremely complicated and has scores of alternatives and subtleties. For example, here are two statements in the forward() method of the TransformerModel class:

positions = self.pos_embedding(T.arange(tk, \
  device=device))[None, :, :].expand(bs, tk, eb)
x = self.unify_embeddings(T.cat((tokens, positions),
  dim=2).view(-1, 2 * eb)).view(bs, tk, eb)

I didn’t conjure up the program from nothing; I used dozens of resources I found on the Internet. There was no single best resource. Transformer architecture is so complicated that different resources were more or less useful to me depending on where I was in my exploration. That said, I found a blog post by Peter Bloem to be exceptionally useful and most of the code comes from that blog post.

Anyway, I’m looking forward to continuing my Transformer architecture exploration.

Three of my favorite fictional book series are set in the British Navy during the 1700s. All involve the main character starting out as a poor midshipman with few prospects, but transforming into a hero through perseverance, courage, and determination. Left: The Horatio Hornblower series was written by C.S. Forester – 10 novels and 5 short stories written from 1937 to 1967. Center: The Jack Aubrey series was written by Patrick O’Brian – 20 novels written from 1969 to 1999. Right: The Richard Bolitho series was written by Douglas Reeman who used the pseudonym Alexander Kent – 30 novels written from 1968 to 2011.

# imdb_transformer.py
# Python 3.7.6, PyTorch 1.7.0
# Windows 10

# ** In Progress - Surely Has Bugs **
# see peterbloem.nl/blog/transformers
# uses BucketIterator - results not reproducible. Arg!
# data has been manually pruned from 25,000 train and
# 25,000 test to 200 train and 200 test

import torch as T
import torchtext as tt
import numpy as np
import math
import warnings  # most of torchtext is deprecated
warnings.filterwarnings("ignore")

device = T.device("cpu")

# -----------------------------------------------------------

class SelfAttention(T.nn.Module):
  def __init__(self, emb_dim, n_heads, mask=False):
    super().__init__()
    self.emb_dim = emb_dim
    self.n_heads = n_heads
    self.mask = mask

    self.tokeys = T.nn.Linear(emb_dim, emb_dim * n_heads, 
      bias=False)
    self.toqueries = T.nn.Linear(emb_dim, emb_dim * n_heads,
      bias=False)
    self.tovalues = T.nn.Linear(emb_dim, emb_dim * n_heads,
      bias=False)
    self.unifyheads = T.nn.Linear(n_heads * emb_dim, 
      emb_dim)

  def forward(self, x):
    bs, tk, eb = x.size()
    nh = self.n_heads
    if eb != self.emb_dim:
      print("Fatal error 01")

    keys    = self.tokeys(x)   .view(bs, tk, nh, eb)
    queries = self.toqueries(x).view(bs, tk, nh, eb)
    values  = self.tovalues(x) .view(bs, tk, nh, eb)

    keys = keys.transpose(1, 2).contiguous().\
      view(bs * nh, tk, eb)
    queries = queries.transpose(1, 2).contiguous().\
      view(bs * nh, tk, eb)
    values = values.transpose(1, 2).contiguous().\
      view(bs * nh, tk, eb)

    # dot product of queries and keys, then scale
    q_k = T.bmm(queries, keys.transpose(1, 2))
    q_k = q_k / math.sqrt(eb) # q_k has attention logits
    if q_k.size() != (bs*nh, tk, tk):
      print("Fatal error 02")

    if self.mask: 
      mask_(q_k, maskval=float('-inf'), \
        mask_diagonal=False)

    q_k = T.nn.functional.softmax(q_k, dim=2)
    # only the first row may contain nan
    if contains_nan(q_k[:, 1:, :]) == True:
      print("Fatal error 03")

    if self.mask == 'first':
      q_k = q_k.clone()
      q_k[:, :1, :] = 0.0
      # The first row of the first attention matrix
      # is entirely masked out, so the softmax operation
      # results in division by zero. Set this row to zero
      # by hand to get rid of the NaNs. PB.

    oupt = T.bmm(q_k, values).view(bs, nh, tk, eb)
    oupt = oupt.transpose(1, 2).contiguous().\
      view(bs, tk, nh * eb)

    return self.unifyheads(oupt)

# -----------------------------------------------------------

class TransformerBlock(T.nn.Module):
  def __init__(self, emb_dim, n_heads, mask, seq_length,
    ff_hid=4, dropout=0.0):
    super().__init__()
    self.attention = SelfAttention(emb_dim, n_heads=n_heads,
      mask=mask)
    self.mask = mask  # False for classification
    self.norm1 = T.nn.LayerNorm(emb_dim)
    self.norm2 = T.nn.LayerNorm(emb_dim)

    self.ff = T.nn.Sequential(
      T.nn.Linear(emb_dim, ff_hid * emb_dim),
      T.nn.ReLU(),
      T.nn.Linear(ff_hid * emb_dim, emb_dim)
    )
    self.do = T.nn.Dropout(dropout)

  def forward(self, x):
    att = self.attention(x)
    x = self.norm1(att + x)
    x = self.do(x)
    f = self.ff(x)
    x = self.norm2(f + x)
    x = self.do(x)

    return x

# -----------------------------------------------------------

class TransformerModel(T.nn.Module):
  def __init__(self, emb_dim, n_heads, depth, seq_length,
    num_tokens, n_classes, max_pool=True, dropout=0.0):

    super().__init__()
    (self.num_tokens, self.max_pool) = num_tokens, max_pool
    self.token_embedding = T.nn.Embedding(embedding_dim=emb_dim,
      num_embeddings=num_tokens)
    self.pos_embedding = T.nn.Embedding(embedding_dim=emb_dim,
      num_embeddings=seq_length)
    self.unify_embeddings = T.nn.Linear(2 * emb_dim, emb_dim)

    tblocks = []
    for i in range(depth):
      tblocks.append(
        TransformerBlock(emb_dim=emb_dim, n_heads=n_heads, \
          seq_length=seq_length, mask=False, dropout=dropout))
    self.tblocks = T.nn.Sequential(*tblocks)

    self.tologits = T.nn.Linear(emb_dim, n_classes)
    self.do = T.nn.Dropout(dropout)

  def forward(self, x):
    tokens = self.token_embedding(x)
    bs, tk, eb = tokens.size()
    positions = self.pos_embedding(T.arange(tk, \
      device=device))[None, :, :].expand(bs, tk, eb)
    x = self.unify_embeddings(T.cat((tokens, positions),
      dim=2).view(-1, 2 * eb)).view(bs, tk, eb)
    x = self.do(x)
    x = self.tblocks(x)
    
    if self.max_pool == True:
      x = x.max(dim=1)[0]
    else:
      x = x.mean(dim=1)
    x = self.tologits(x)

    return T.nn.functional.log_softmax(x, dim=1)

# -----------------------------------------------------------

def contains_nan(tnsr):
  return bool((tnsr != tnsr).sum() > 0)

# -----------------------------------------------------------

def predict(model, review, TXT, max_len):
  # assumes model.eval()
  # review = "this movie was not a waste of my time ."
  lst = []
  words = review.split()

  n = len(words)
  for i in range(n):
    id = TXT.vocab.stoi[words[i]]
    lst.append(id)
  for i in range(max_len-n):
    lst.append(1)  # padding

  inpt = T.tensor(lst, dtype=T.int64).to(device)
  inpt = inpt.reshape(1, -1)
  with T.no_grad():
    oupt = model(inpt)

  result = T.exp(oupt)
  return result

# -----------------------------------------------------------

def main():
  print("\nBegin IMDB example using from-scratch Transformer ")

  # 0. get ready
  T.manual_seed(1)
  np.random.seed(1)

  # 1. set up parameters

  # 1a. data
  arg_batch_size = 50
  arg_max_vocab_words = 3_500  # 50_000

  # 1b. model
  arg_embedding_size = 128  # 128
  arg_num_heads = 8  # 8
  arg_depth = 6  # 6
  arg_max_seq_len  = 60  # 512
  arg_num_cls = 2
  arg_max_pool = True

  # 1c. training
  arg_lr = 0.01  # 0.0001
  arg_lr_warmup = 100  # 10_000 
  arg_num_epochs = 3  # 80
  arg_gradient_clip = 1.0  # 1.0

  # 2. create IMDB data iterators
  print("\nGetting data and building vocabulary ")
  TEXT = tt.data.Field(lower=True, include_lengths=True,
    batch_first=True, tokenize="basic_english")
  LABEL = tt.data.Field(sequential=False)

  train_ds, test_ds = tt.datasets.IMDB.splits(TEXT, LABEL)
  TEXT.build_vocab(train_ds, max_size=arg_max_vocab_words-2) 
  LABEL.build_vocab(train_ds)  # unk neg, pos

  train_itr, test_itr = \
    tt.data.BucketIterator.splits((train_ds, test_ds), \
    shuffle=True, batch_size=arg_batch_size)

# -----------------------------------------------------------
 
  # 3. create model
  print("\nCreating Transformer binary classifier model ")
  model = TransformerModel(emb_dim=arg_embedding_size,
    n_heads=arg_num_heads, depth=arg_depth,
    seq_length=arg_max_seq_len, num_tokens=arg_max_vocab_words,
    n_classes=arg_num_cls, max_pool=arg_max_pool).to(device)
  print("Model created ")

  # 4. train model
  opt = T.optim.Adam(lr=arg_lr, params=model.parameters())
  sch = T.optim.lr_scheduler.LambdaLR(opt, lambda i: \
    min(i / (arg_lr_warmup / arg_batch_size), 1.0))

  for epoch in range(arg_num_epochs):
    print("\nEpoch " + str(epoch))
    model.train(True)

    for bat_idx, batch in enumerate(train_itr):
      opt.zero_grad()

      inpt = batch.text[0].to(device)
      # lengths = batch.text[1]
      label = (batch.label - 1).to(device)  # 0, 1, 2 

      if inpt.size(1) "greater-than" arg_max_seq_len: # replace
        inpt = inpt[:, :arg_max_seq_len]
      oupt = model(inpt)  # like [-0.123, -0.456] (logits)

      loss = T.nn.functional.nll_loss(oupt, label) 
      loss.backward()

      # clip gradients
      if arg_gradient_clip "greater-than" 0.0: # replace
        T.nn.utils.clip_grad_norm_(model.parameters(), 
          arg_gradient_clip)

      opt.step()  # update weights and biases
      sch.step()  # adjust the LR

      if bat_idx % 1 == 0:
        print(" batch: " + str(bat_idx), end="")
        print("  loss = %0.4f " % loss)

    with T.no_grad():  # after each epoch
      print("epoch finished. TODO: compute accuracy") 
      model.eval()
      # TODO: overall loss, test loss
      # TODO: save checkpoint

  # 5. use model to make prediction
  review = "this movie was not a waste of my time ."
  print("\nUsing model to predict :")
  print("\"" + review + "\"")
  probs = predict(model, review, TEXT, arg_max_seq_len)
  print("\nPredicted probs [neg, pos]: ")
  print(probs)

  # 6. TODO: compute train, test accuracy

  # 7. TODO: save trained model

# -----------------------------------------------------------

if __name__ == "__main__":
  main()