Hyperparameters from papers

src/models/train_model.py

@@ -10,15 +10,25 @@ from src.models.predict_model import evaluate
 
 def get_model(name: str, triples, inverse_triples):
     model_dict = {
-        "TransE": TransE(triples_factory=triples),
-        "TransR": TransR(triples_factory=triples),
-        "ConvE": ConvE(triples_factory=inverse_triples),
-        "DistMult": DistMult(triples_factory=triples),
+        "TransE": TransE(triples_factory=triples, embedding_dim=20,
+                         scoring_fct_norm=1),
+        "TransR": TransR(triples_factory=triples, embedding_dim=50,
+                         relation_dim=50, scoring_fct_norm=1),
+        "ConvE": ConvE(triples_factory=inverse_triples, embedding_dim=200,
+                       feature_map_dropout= 0.2,  input_dropout= 0.2,
+                       output_dropout=0.3),
+        "DistMult": DistMult(triples_factory=triples, embedding_dim=100),
     }
-    return model_dict[name]
+    lr_dict = {
+        "TransE": 0.01,
+        "TransR": 0.001,
+        "ConvE": 0.001,
+        "DistMult": 0.1,
+    }
+    return model_dict[name], lr_dict[name]
 
 
-def train(model, train_data, optimizer, n_epochs: int = 5, batch_size: int = 256):
+def train(model, train_data, optimizer, n_epochs: int = 20, batch_size: int = 256):
     training_loop = SLCWATrainingLoop(
         model=model,
         triples_factory=train_data,
@@ -43,9 +53,9 @@ if __name__ == '__main__':
     for name in model_names:
         model_results[name] = {}
         print(name)
-        model = get_model(name, training_triples_factory, training_inverse)
+        model, lr = get_model(name, training_triples_factory, training_inverse)
 
-        optimizer = Adam(params=model.get_grad_params())
+        optimizer = Adam(lr=lr, params=model.get_grad_params())
 
         train(model, training_triples_factory, optimizer)
 
@@ -64,7 +74,7 @@ if __name__ == '__main__':
         model_results[name]["MR"] = results.data[('arithmetic_mean_rank', 'both', 'realistic')]
         model_results[name]["MRR"] = results.data[('inverse_harmonic_mean_rank', 'both', 'realistic')]
 
-    for name in model_names:
+    # for name in model_names:
         print(name)
         print(f"Hits@1: {model_results[name]['Hits@1']}")
         print(f"Hits@3: {model_results[name]['Hits@3']}")