From d912e0a2251ca5911d27ca0c800cbb8a491dc8b6 Mon Sep 17 00:00:00 2001 From: thanhvc3 Date: Wed, 17 May 2023 13:41:13 +0700 Subject: [PATCH] add grid search --- main.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/main.py b/main.py index 8986c2d..d1487c5 100644 --- a/main.py +++ b/main.py @@ -669,8 +669,13 @@ if __name__ == "__main__": for step, batch in enumerate(dataloader): sub, rel, obj, label, neg_ent, sub_samp = model.read_batch( batch, 'train') + + if (neg_ent is None): + neg_ent = np.repeat(None, repeats=len(sub)) + else: + neg_ent = neg_ent.cpu() - dataset = np.stack([sub.cpu(), rel.cpu(), neg_ent.cpu(), np.repeat(model.p.train_strategy, repeats=len(sub))], axis = 1) + dataset = np.stack([sub.cpu(), rel.cpu(), neg_ent, np.repeat(model.p.train_strategy, repeats=len(sub))], axis = 1) search = grid.fit(dataset, label) print("BEST SCORE: ", search.best_score_) print("BEST PARAMS: ", search.best_params_)