def convert_trax_checkpoint_to_pytorch(trax_model_pkl_path, config_file, pytorch_dump_path):
# Initialise PyTorch model
config = ReformerConfig.from_json_file(config_file)
print("Building PyTorch model from configuration: {}".format(str(config)))
model = ReformerModelWithLMHead(config)
with open(trax_model_pkl_path, "rb") as f:
model_weights = pickle.load(f)["weights"]
set_model_weights_in_torch(model_weights, model, config.hidden_size)
# Save pytorch-model
print("Save PyTorch model to {}".format(pytorch_dump_path))
torch.save(model.state_dict(), pytorch_dump_path)
with torch.no_grad():
tmp = next(val_loader)
input_ids = tmp['input_ids']
attention_mask = tmp['attention_mask']
labels = tmp['labels']
outputs = model(input_ids,
attention_mask=attention_mask,
labels=labels)
loss, prediction_scores = outputs[:2]
val_loss[f"Epoch {x} Step {i}"] = loss.item()
all_val_loss[f"Epoch {x} Step {i}"] = loss.item()
print(f'validation loss: {loss.item()}')
torch.save(model.state_dict(),
f"saved_0624/model/saved_model_epoch_{x}.pth")
with open(f'saved_0624/saved_losses/training_loss_epoch_{x}.json',
'w') as f:
f.write(json.dumps(training_loss))
with open(f'saved_0624/saved_losses/val_loss_epoch_{x}.json', 'w') as f:
f.write(json.dumps(val_loss))
end = time.time()
print(f"----------{(end-start)//60} min per epoch----------")
with open("saved_0624/saved_losses/training_loss_all.json", 'w') as f:
f.write(json.dumps(all_training_loss))
with open("saved_0624/saved_losses/val_loss_all.json", 'w') as f:
f.write(json.dumps(all_val_loss))
