How to save and continue training for skipgram negative sampling

username: vortana_say
cooked: <p>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.</p>
<p>Example training code:</p>
<pre><code class="lang-auto">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")
</code></pre>
<p>I am looking to the way to save and continue training efficiently.</p>
<p>Thanks.</p>

username: ptrblck
cooked: <p>You could adapt the <code>save_checkpoint</code> method from the <a href="https://github.com/pytorch/examples/blob/391be73bbd191a08776c9ba05d29ef1e05c0c83b/imagenet/main.py#L254-L260">ImageNet example</a> to store a checkpoint for each epoch.<br>
Note that in order to continue the training from a checkpoint you should at least store the <code>state_dict</code>s of the model as well as the optimizer.</p>

username: vortana_say
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> 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:</p>
<pre><code class="lang-auto">checkpoint = {
            'epoch': epoch + 1,
            'valid_loss_min': valid_loss,
            'state_dict': model.state_dict(),
            'optimizer': optimizer.state_dict(),
        }
</code></pre>
<p>For NLP, are there any specific thing I should save on top of the things above ?</p>

username: ptrblck
cooked: <p>You could add everything into the <code>dict</code>, which would be needed to continue the training.<br>
E.g. the setup of a learning rate scheduler, model and preprocessing arguments etc.</p>
<p>I would also recommend to run a quick test for a few epochs, store the <code>checkpoint</code>, and make sure the continuation of the training is working well.</p>

username: vortana_say
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> I will try that out, and update. Thanks</p>