| 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') | |