How to save and continue training for skipgram negative sampling

nlp
1

Hi, I am new with skipgram. I went through the documents, it is taking a lot of time to train both skipgram and / or skipgram with negative sampling. The issue happens when the notebook connection is disconnected.

Example training code:

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

# Get our noise distribution
# Using word frequencies calculated earlier in the notebook
word_freqs = np.array(sorted(freqs.values(), reverse=True))
unigram_dist = word_freqs/word_freqs.sum()
noise_dist = torch.from_numpy(unigram_dist**(0.75)/np.sum(unigram_dist**(0.75)))

# instantiating the model
embedding_dim = 300
model = SkipGramNeg(len(vocab_to_int), embedding_dim, noise_dist=noise_dist).to(device)

# using the loss that we defined
criterion = NegativeSamplingLoss() 
optimizer = optim.Adam(model.parameters(), lr=0.003)

print_every = 1500
steps = 0
epochs = 5

# train for some number of epochs
for e in range(epochs):
    
    # get our input, target batches
    for input_words, target_words in get_batches(train_words, 512):
        steps += 1
        inputs, targets = torch.LongTensor(input_words), torch.LongTensor(target_words)
        inputs, targets = inputs.to(device), targets.to(device)
        
        # input, outpt, and noise vectors
        input_vectors = model.forward_input(inputs)
        output_vectors = model.forward_output(targets)
        noise_vectors = model.forward_noise(inputs.shape[0], 5)

        # negative sampling loss
        loss = criterion(input_vectors, output_vectors, noise_vectors)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # loss stats
        if steps % print_every == 0:
            print("Epoch: {}/{}".format(e+1, epochs))
            print("Loss: ", loss.item()) # avg batch loss at this point in training
            valid_examples, valid_similarities = cosine_similarity(model.in_embed, device=device)
            _, closest_idxs = valid_similarities.topk(6)

            valid_examples, closest_idxs = valid_examples.to('cpu'), closest_idxs.to('cpu')
            for ii, valid_idx in enumerate(valid_examples):
                closest_words = [int_to_vocab[idx.item()] for idx in closest_idxs[ii]][1:]
                print(int_to_vocab[valid_idx.item()] + " | " + ', '.join(closest_words))
            print("...\n")

I am looking to the way to save and continue training efficiently.

Thanks.

2

You could adapt the save_checkpoint method from the ImageNet example to store a checkpoint for each epoch.
Note that in order to continue the training from a checkpoint you should at least store the state_dicts of the model as well as the optimizer.

3

@ptrblck I will try that; I think I am confused because I have been more accustomed to train the image; and I save latest checkpoint and best checkpoints:

checkpoint = {
            'epoch': epoch + 1,
            'valid_loss_min': valid_loss,
            'state_dict': model.state_dict(),
            'optimizer': optimizer.state_dict(),
        }

For NLP, are there any specific thing I should save on top of the things above ?

4

You could add everything into the dict, which would be needed to continue the training.
E.g. the setup of a learning rate scheduler, model and preprocessing arguments etc.

I would also recommend to run a quick test for a few epochs, store the checkpoint, and make sure the continuation of the training is working well.

5

@ptrblck I will try that out, and update. Thanks