ML: Recurrent Networks in PyTorch

---

Introduction

---

Encoder-Decoder

Suppose we have two different recurrent networks. The first is called Encoder, and the second is Decoder. The encoder receives text in some language (source) as input, while the decoder should output text in another language (target) at its output. The last cell of the RNN encoder contains at its output the hidden state vector $\mathbf{h}_e^{\text{last}}$ of its last cell. This vector "stores" information about the entire source sentence (the context vector). It is fed as the initial hidden state to the first cell of the RNN decoder.

Then the service token <BOS> (begin of sentence) is fed to the input of the first decoder cell. The word-translation "кот" should appear at the cell output. This means that the output of each cell is passed through a linear layer with the number of neurons equal to the number of words in the vocabulary. Then the softmax function outputs word probabilities, from which the index of the maximum is selected (argmax). The resulting word "кот" is passed to the input of the second cell and so on until the service token <EOS> (end of sentence) is obtained.

A faster but less accurate training mode is called teacher forcing. In this case, the correct target translation is fed to all decoder inputs at once, and the decoder is required to output this sentence shifted one word to the left. Usually, random switching occurs between the "honest" and "forced" training modes.

---

Encoder-Decoder implementation

Here is an implementation of the Encoder-Decoder architecture in PyTorch. Let the source language vocabulary size be VOC_SIZE and the embedding vector dimension E = VEC_DIM. Then the encoder module looks like this:

VEC_DIM = 100

class EncoderRNN(nn.Module):
    def __init__(self, VOC_SIZE, E):                  # vocabulary size & embedding dimension
        super(EncoderRNN, self).__init__()
 
        self.emb = nn.Embedding(VOC_SIZE, E, scale_grad_by_freq=True)      
        self.rnn = nn.GRU(E, E, bidirectional=True)   # bidirectional GRU   
  
    def forward(self, X):                               
        """ X:(B,L)  B - number of sentences with L words each. Zero means absence of a word """
         
        lens = torch.tensor([ len(x)-len(x[x==0]) for x in X ])
         
        emb = self.emb( X.t() )                       # (B,L) -> (L,B) -> (L,B,E)     
         
        Xp   = pack_padded_sequence(emb, lens, enforce_sorted=False)        
                  
        _, Hn = self.rnn(Xp)                          # (2,B,E)
        Hn = torch.cat([Hn[0],Hn[1]], dim=1)          # (B,2*E)
         
        return Hn.view(1,-1, Hn.size(1))              # (1,B,2*E) only hidden state

We will send one sentence at a time to the encoder (batch_size=1) of variable length L in the form of a vector of L integers (long). At the output, the encoder returns a pair: Y - tensor (L,1,E) of all cell outputs and the output Hid of the last cell.

class DecoderRNN(nn.Module):
    def __init__(self, VOC_SIZE, E):
        super(DecoderRNN, self).__init__()

        self.emb = nn.Embedding(VOC_SIZE, E, scale_grad_by_freq=True)
        self.rnn = nn.GRU(E, 2*E)                    # hidden from bidirectional
        self.out = nn.Linear(2*E, VOC_SIZE)        

    def forward(self, Hid, X = None, forcing = False):     #  Hid:(1,B,2*E), X:(B,L)
        max_len = MAX_EN_LEN if X is None else len(X[0])   # maximum sentence length
                
        W    = torch.empty( Hid.size(1),           dtype=torch.long ).fill_(BOS_INDEX)
        Wrds = torch.zeros( Hid.size(1),  max_len, dtype=torch.long  ) # predicted words
        Prbs = torch.ones ( Hid.size(1),  max_len, dtype=torch.float )  # probabilities 
        
        for i in range(max_len):                           #                
            W = self.emb( W.view(1,-1) )                   # (1,B,E)
            Y, Hid = self.rnn(W, Hid)                      # (1,B,2*E)
            
            Y = self.out (Y[0])                            # (B,VOC_SIZE)
            Y = torch.softmax( Y, dim=1 )                  # (B,VOC_SIZE)
            _, W =  Y.detach().topk(1, dim=1)              #  (B,1)            
            
            Wrds[:,i].copy_(W.squeeze())                   # remove 1 and save
            if not X is None and i < X.size(1):
                for b in range(X.size(0)): Prbs[b,i] = Y[b, X[b,i]]
            if forcing:
                W.copy_( X[:,i].view(-1,1) )
                        
        return Wrds, Prbs                                  # (B,L), (B,L)

The combined model:

class EncoderDecoderRNN(nn.Module):
    def __init__(self, encoder, decoder):
        super(EncoderDecoderRNN, self).__init__()
        
        self.enc = encoder
        self.dec = decoder

    def forward(self, sourse, target, forcing=False):           # (L1,)  (L2,)
        Hd = self.enc(sourse)        
        return self.dec(Hd, target, forcing)
                    
gpu = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
cpu = torch.device("cpu")

encoder = EncoderRNN(len(voc_en), VEC_DIM)
decoder = DecoderRNN(len(voc_ru), VEC_DIM)

model = EncoderDecoderRNN(encoder, decoder)                            # model instance 
model.to(gpu) 

---