copy_task.py

import torch
import torch.nn as nn
import dni

BATCH_SIZE = 16
D_MODEL = 256
TBPTT = 3
COPY_SIZE = 3*TBPTT - 1
ALPHABET = "01"
EPOCHS = 10

# TODO: add arg parser
    # make clip grad a param
# TODO: show bits error with more comprehensive output (see pytorch WLM example)


class LSTM_plus_embedding(nn.Module):
    def __init__(self, size):
        super().__init__()
        self.embedding = nn.Embedding(len(ALPHABET)+1, size)
        self.rnn = nn.LSTM(size, hidden_size = size)
        self.output_embedding = nn.Linear(size, len(ALPHABET)+1)
        
        # # weight tying
        # self.output_embedding.weight = self.embedding.weight
       
        
    def forward(self, input, h0):
        x = self.embedding(input)
        out, hn = self.rnn(x, h0)
        out = self.output_embedding(out)
        return out, hn
    

if __name__ == "__main__":
    device = 'cuda'
    
    ### DATA ###
    # create prompts with a stop character at the end
    batched_prompts = torch.randint(0, len(ALPHABET), (EPOCHS, COPY_SIZE, BATCH_SIZE), dtype = torch.long)
    batched_prompts = torch.cat((batched_prompts, len(ALPHABET)*torch.ones(EPOCHS, 1, BATCH_SIZE, dtype = torch.long)), dim=1)
    batched_prompts = torch.utils.data.DataLoader(batched_prompts, batch_size = None, pin_memory = True)
    
    rnn = LSTM_plus_embedding(D_MODEL).to(device)
    optim = torch.optim.SGD(rnn.parameters(), lr = .001)
    loss = nn.CrossEntropyLoss()

    # instantiate DNI model that will backward synthetic gradients
    synth = dni.Synthesizer(size = D_MODEL, is_lstm = True).cuda()
    
    losses = []
    synth_losses = []
    aux_losses = []
    synths = []
    for epoch, batch in enumerate(batched_prompts):
        batch = batch.to(device)
        
        # initialize the hidden state with synth
        h_n = (torch.ones(1, BATCH_SIZE, D_MODEL).to(device), torch.ones(1, BATCH_SIZE, D_MODEL).to(device))

        # split into TBPTT size sections
        for split in torch.split(batch, TBPTT, dim = 0):
            # standard forward pass
            out, h_n = rnn(split, h_n)
            cross_loss = loss(out.view(-1, len(ALPHABET)+1), torch.zeros(BATCH_SIZE, TBPTT, dtype = torch.long).view(-1).to(device))
            cross_loss.backward(retain_graph = True)
            
            # just add ONE line for synthetic gradients
            h_n = synth.backward_synthetic(h_n)

            torch.nn.utils.clip_grad_norm_(rnn.parameters(), 25)
            optim.step()
            optim.zero_grad()

        # check copy ability
        for split in torch.split(batch, TBPTT, dim = 0):
            # standard forward pass
            out, h_n = rnn(torch.zeros(TBPTT, BATCH_SIZE, dtype = torch.long).to(device), h_n)
            cross_loss = loss(out.reshape(-1, len(ALPHABET)+1), split.reshape(-1))
            cross_loss.backward(retain_graph = True)
            
            # just add ONE line for synthetic gradients
            h_n = synth.backward_synthetic(h_n)

            losses.append(cross_loss.item())

            torch.nn.utils.clip_grad_norm_(rnn.parameters(), 25)
            optim.step()
            optim.zero_grad()
        
        # step the the DNI model 
        synth.step()

        print(f'epoch: {epoch}', end = '\r')