Skip to content
Snippets Groups Projects
Commit 29060a91 authored by lenhy's avatar lenhy
Browse files

Margin in TransE loss

parent 4d37a7fc
Branches
No related tags found
No related merge requests found
......@@ -45,7 +45,10 @@ def get_model(name: str, triples, inverse_triples, regularizer):
"ConvE": 0.1,
"DistMult": 0.0,
}
return model_dict[name], lr_dict[name], epochs_dict[name], bs_dict[name], ls_dict[name]
if name == "TransE":
model = model_dict[name]
model.loss.margin = 2.0
return model, lr_dict[name], epochs_dict[name], bs_dict[name], ls_dict[name]
def train(model, train_data, eval_data, optimizer, evaluator, n_epochs: int = 20, batch_size: int = 256, label_smoothing: float = 0.0):
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment