File size: 1,160 Bytes
f9c98cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 | import torch
import numpy as np
import sys
def save_embed_npz(dataset, model_path, data_path, model_name):
sys.path.append('./' + str(model_path))
experiment_path = model_path / 'experiments' / (dataset + '_' + model_name) / 'best_model.pth'
save_path = data_path.parent / 'embed_model' / (dataset + '_' + model_name + '_embed.npz')
model = torch.load(experiment_path)
if model_name == 'ComplEx':
rel_embeddings = torch.sqrt(torch.square(model.rel_re_embeddings.weight.data) + torch.square(model.rel_im_embeddings.weight.data))
ent_embeddings = torch.sqrt(torch.square(model.ent_re_embeddings.weight.data) + torch.square(model.ent_im_embeddings.weight.data))
else:
rel_embeddings = model.rel_embeddings.weight.data
ent_embeddings = model.ent_embeddings.weight.data
eM = np.array(ent_embeddings.cpu())
rM = np.array(rel_embeddings.cpu())
if not save_path.parent.exists():
save_path.parent.mkdir()
np.savez(save_path, eM=eM, rM=rM)
print(model_name+' saved successfully！！！')
if __name__ == '__main__':
save_embed_npz('wordnet', 'TransE')
|