import torch
from model import MainModel
from dataloader import testDataset
from tqdm import tqdm
import numpy as np
from sklearn.metrics import f1_score
dataset = testDataset()
model = MainModel(10310, 15439).cuda()
model.load_state_dict(torch.load("model_lstm.pth"))
total = 0
correct = 0
outs = []
preds = []
with torch.no_grad():
for idx, data in tqdm(enumerate(dataset)):
outputs = model(data["source"].unsqueeze(0).cuda(),
data["target"].unsqueeze(0).cuda(),
data["alignment"].unsqueeze(0).cuda()).view(-1).cpu()
# print(outputs.numpy() > 0.5)
ans = np.sum((outputs.numpy() > 0.5) == data["predictions"].numpy())
outs.extend((outputs.numpy() > 0.5))
preds.extend(data["predictions"].numpy())
# print(ans, outputs.size()[0])
total += outputs.size()[0]
correct += ans
train_data_num = ceil(cascade_num / 10 * 9)
syn_small_marker = torch.load("data/" + dataset +
"/marker.pkl").cuda(0)[train_data_num:, :]
syn_small_time = torch.load("data/" + dataset +
"/time.pkl").cuda(0)[train_data_num:, :]
syn_small_mask = torch.load("data/" + dataset +
"/mask.pkl").cuda(0)[train_data_num:, :]
if small:
syn_small_adj_mat = torch.load("data/" + dataset + "/adj_mat.pkl")
syn_small_adj_list = torch.load("data/" + dataset + "/adj_list.pkl")
syn_small_main = MainModel(marker_num, 5, d_model, d_model, 2 * d_model,
d_model, d_model, d_model, head, candi_size,
max_time, 0.3, 0, 10, 0.99, 0.001, 0.1)
syn_small_main.load_state_dict(torch.load("models/main-" + dataset +
"-512.pt"))
syn_small_main.cpu().eval()
if small:
syn_small_evaluator = Small_Network_Evaluator(marker_num, neg_num,
test_size, syn_small_adj_mat)
else:
syn_small_evaluator = Large_Network_Evaluator(marker_num, neg_num,
test_size,
syn_small_adj_list)
main = []
main_mean = []
for i in range(3):
main_i = []