def load_checkpoint(filename, dir=GOOGLE_COLAB_CHECKPOINT_DIR):
checkpoint = torch.load(os.path.join(dir, filename))
model = RNN(hidden_size=checkpoint['hidden_size'],
output_size=checkpoint['output_size'],
n_layers=checkpoint['n_layers'],
batch_size=checkpoint['batch_size'],
bidirectional=checkpoint['bidirectional'])
model.load_state_dict(checkpoint['state_dict'])
model.num_epochs_trained = checkpoint['num_epochs_trained']
return model
import os
import dataset
from network import RNN
dataset = dataset.Dataset()
max_length = 20
rnn = RNN(dataset.n_letters, 128, dataset.n_letters, dataset.n_categories)
rnn.eval()
# load weights
if os.path.exists("models/gen_names.pkl"):
checkpoint = torch.load("models/gen_names.pkl")
rnn.load_state_dict(checkpoint['nn_state_dict'])
print("checkpoint loaded")
def sample(category, start_char):
with torch.no_grad():
category_tensor_var = dataset.category_tensor(category)
input = dataset.input_tensor(start_char)
hidden = rnn.init_hidden()
output_name = start_char
for i in range(max_length):
output, hidden = rnn(category_tensor_var, input[0], hidden)
topv, topi = output.topk(1)
topi = topi[0][0]
if topi == dataset.n_letters - 1: