Get predictions to actually work

Changes

2 changed files (+20/-17)

• predict.py
  • @@ -2,16 +2,16 @@ import numpy as np
    import torch - def predict(device, dataset, model, text, next_words=100, top_k=5): + def predict(device, dataset, model, text, next_words=100, top_k=3): model.eval() - words = text.split(' ') + words = text.split() state_h, state_c = model.zero_state(1) state_h = state_h.to(device) state_c = state_c.to(device) - for w in words: + for word in words: ix = torch.tensor([[dataset.word_to_index[word]]]).to(device) output, (state_h, state_c) = model(ix, (state_h, state_c))
  • @@ -23,6 +23,6 @@ def predict(device, dataset, model, text, next_words=100, top_k=5):
    words.append(dataset.index_to_word[choice]) ix = torch.tensor([[choice]]).to(device) - output, (state_h, state_c) = net(ix, (state_h, state_c)) + output, (state_h, state_c) = model(ix, (state_h, state_c)) return words
• train.py
  • @@ -6,6 +6,7 @@ from torch.utils.data import DataLoader
    from dataset import Dataset from model import Model + from predict import predict parser = ArgumentParser()
  • @@ -19,9 +20,9 @@ parser.add_argument('-s', '--seq-size', default=32, type=int,
    help='sequence size') parser.add_argument('-b', '--batch-size', default=256, type=int, help='size of each training batch') - parser.add_argument('-m', '--embedding-dim', default=512, type=int, + parser.add_argument('-m', '--embedding-dim', default=256, type=int, help='size of the embedding') - parser.add_argument('-l', '--lstm-size', default=512, type=int, + parser.add_argument('-l', '--lstm-size', default=256, type=int, help='size of the LSTM hidden state') parser.add_argument('-a', '--layers', default=3, type=int, help='number of LSTM layers')
  • @@ -39,8 +40,9 @@ print(len(dataloader))
    # Prepare model + device = torch.device(args.device) model = Model(dataset, args.embedding_dim, args.lstm_size, - args.layers, args.dropout).to(args.device) + args.layers, args.dropout).to(device) print(model)
  • @@ -50,20 +52,21 @@ optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    for t in range(args.epochs): state_h, state_c = model.zero_state(args.batch_size) - state_h = state_h.to(args.device) - state_c = state_c.to(args.device) + state_h = state_h.to(device) + state_c = state_c.to(device) iteration = 0 + print(len(dataloader)) for batch, (X, y) in enumerate(dataloader): - model.train() - iteration += 1 + model.train() + optimizer.zero_grad() - X = torch.tensor(X).to(args.device) - y = torch.tensor(y).to(args.device) + X = torch.tensor(X).to(device) + y = torch.tensor(y).to(device) # Compute prediction error logits, (state_h, state_c) = model(X, (state_h, state_c))
  • @@ -87,7 +90,7 @@ for t in range(args.epochs):
    'Iteration: {}'.format(iteration), 'Loss: {}'.format(loss_value)) - if iteration % 1000 == 0: - predict(args.device, dataset, model, 100, 2) - torch.save(net.state_dict(), - 'checkpoint/model-{}.pth'.format(iteration)) + if iteration % 1 == 0: + print(' '.join(predict(args.device, dataset, model, 'i am'))) + #torch.save(model.state_dict(), + # 'checkpoint/model-{}.pth'.format(iteration))