def load_model(config): # todo
if type(config) == str:
with open(config) as file:
content = file.read()
model_dir = os.path.abspath(os.path.join(config, os.pardir))
else:
model_dir = None
with open(os.path.join(model_dir, "configuration.json")) as f:
config = json.load(f)
kge_model = pipeline.get_kge_model(config=config)
if model_dir is not None:
path_to_model = os.path.join(model_dir, "trained_model.pkl")
kge_model.load_state_dict(torch.load(path_to_model))
return kge_model
def train(self, pos_triples: np.array, neg_triples: np.array):
all_triples = np.concatenate((pos_triples, neg_triples))
# testme
self.config[keenConst.NUM_ENTITIES] = len(
np.unique(np.concatenate((all_triples[:, 0], all_triples[:, 2]))))
self.config[keenConst.NUM_RELATIONS] = len(np.unique(all_triples[:,
1]))
## prepare model
self.kge_model = pipeline.get_kge_model(config=self.config)
self.kge_model = self.kge_model.to(self.device)
optimizer = optim.SGD(self.kge_model.parameters(),
lr=self.config[keenConst.LEARNING_RATE])
loss_per_epoch = []
num_pos_examples = pos_triples.shape[0]
num_neg_examples = neg_triples.shape[0]
### train model
tqdmbuffer = TqdmBuffer() if globConst.GUI_MODE else None
for _ in tqdm(range(self.config[keenConst.NUM_EPOCHS]),
file=tqdmbuffer):
# create batches
indices_pos = np.arange(num_pos_examples)
np.random.shuffle(indices_pos)
pos_triples = pos_triples[indices_pos]
pos_batches = self._split_list_in_batches(
input_list=pos_triples, batch_size=self.config["batch_size"])
indices_neg = np.arange(num_neg_examples)
np.random.shuffle(indices_neg)
neg_triples = neg_triples[indices_neg]
neg_batches = self._split_list_in_batches(
input_list=neg_triples, batch_size=self.config["batch_size"])
current_epoch_loss = 0.0
#tqdmbuffer = TqdmBuffer() if globConst.GUI_MODE else None
for pos_batch, neg_batch in tqdm(zip(pos_batches, neg_batches),
total=len(neg_batches)):
current_batch_size = len(pos_batch)
# if not len(pos_batch) == len(neg_batch):
# raise RuntimeError('Pos and neg batches are not the same length')
pos_batch_tensor = torch.tensor(pos_batch,
dtype=torch.long,
device=self.device)
neg_batch_tensor = torch.tensor(neg_batch,
dtype=torch.long,
device=self.device)
# Recall that torch *accumulates* gradients. Before passing in a
# new instance, you need to zero out the gradients from the old instance
optimizer.zero_grad()
loss = self.kge_model(pos_batch_tensor, neg_batch_tensor)
current_epoch_loss += loss.item() * current_batch_size
loss.backward()
optimizer.step()
loss_per_epoch.append(current_epoch_loss / len(pos_triples))
### prepare results for output
entity_to_embedding = {
id: embedding.detach().cpu().numpy()
for id, embedding in enumerate(
self.kge_model.entity_embeddings.weight)
}
relation_to_embedding = {
id: embedding.detach().cpu().numpy()
for id, embedding in enumerate(
self.kge_model.relation_embeddings.weight)
}
results = {
"trained_model": self.kge_model,
"loss_per_epoch": loss_per_epoch,
"entity_to_embedding": entity_to_embedding,
"relation_to_embedding": relation_to_embedding,
"config": self.config,
}
self.output_results(results)
return self.kge_model