def prepare_data(args, graph, target_type, train_target_nodes, valid_target_nodes, pool):
"""
Sampled and prepare training and validation data using multi-process parallization.
"""
jobs = []
for batch_id in np.arange(args.n_batch):
p = pool.apply_async(
node_classification_sample,
args=(args, target_type, randint(), train_target_nodes, {1: True}),
)
jobs.append(p)
p = pool.apply_async(
node_classification_sample,
args=(args, target_type, randint(), valid_target_nodes, {1: True}),
)
jobs.append(p)
return jobs
def fit(
self, model: SupervisedHeterogeneousNodeClassificationModel, dataset: Dataset
) -> None:
args = self.args
self.device = args.device_id[0] if not args.cpu else "cpu"
self.data = preprocess_dataset(dataset)
global graph_pool
graph_pool = self.data
self.target_type = "def"
self.train_target_nodes = self.data.train_target_nodes
self.valid_target_nodes = self.data.valid_target_nodes
self.test_target_nodes = self.data.test_target_nodes
self.types = self.data.get_types()
self.criterion = torch.nn.NLLLoss()
self.stats = []
self.res = []
self.best_val = 0
self.train_step = 0
self.pool = mp.Pool(args.n_pool)
self.st = time.time()
self.jobs = prepare_data(
args,
self.data,
self.target_type,
self.train_target_nodes,
self.valid_target_nodes,
self.pool,
)
"""
Initialize GNN (model is specified by conv_name) and Classifier
"""
self.gnn = GNN(
conv_name=args.conv_name,
in_dim=len(self.data.node_feature[self.target_type]["emb"].values[0]),
n_hid=args.n_hid,
n_heads=args.n_heads,
n_layers=args.n_layers,
dropout=args.dropout,
num_types=len(self.types),
num_relations=len(self.data.get_meta_graph()) + 1,
prev_norm=args.prev_norm,
last_norm=args.last_norm,
use_RTE=False,
)
if args.use_pretrain:
self.gnn.load_state_dict(
load_gnn(torch.load(args.pretrain_model_dir)), strict=False
)
print("Load Pre-trained Model from (%s)" % args.pretrain_model_dir)
self.classifier = Classifier(args.n_hid, self.data.y.max().item() + 1)
self.model = torch.nn.Sequential(self.gnn, self.classifier).to(self.device)
self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=5e-4)
if args.scheduler == "cycle":
self.scheduler = torch.optim.lr_scheduler.OneCycleLR(
self.optimizer,
pct_start=0.02,
anneal_strategy="linear",
final_div_factor=100,
max_lr=args.max_lr,
total_steps=args.n_batch * args.n_epoch + 1,
)
elif args.scheduler == "cosine":
self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
self.optimizer, 500, eta_min=1e-6
)
else:
assert False
self.train_data = [job.get() for job in self.jobs[:-1]]
self.valid_data = self.jobs[-1].get()
self.pool.close()
self.pool.join()
self.et = time.time()
print("Data Preparation: %.1fs" % (self.et - self.st))
for epoch in np.arange(self.args.n_epoch) + 1:
"""
Prepare Training and Validation Data
"""
train_data = [job.get() for job in self.jobs[:-1]]
valid_data = self.jobs[-1].get()
self.pool.close()
self.pool.join()
"""
After the data is collected, close the pool and then reopen it.
"""
self.pool = mp.Pool(self.args.n_pool)
self.jobs = prepare_data(
self.args,
self.data,
self.target_type,
self.train_target_nodes,
self.valid_target_nodes,
self.pool,
)
self.et = time.time()
print("Data Preparation: %.1fs" % (self.et - self.st))
"""
Train
"""
self.model.train()
train_losses = []
for (
node_feature,
node_type,
edge_time,
edge_index,
edge_type,
x_ids,
ylabel,
) in train_data:
node_rep = self.gnn.forward(
node_feature.to(self.device),
node_type.to(self.device),
edge_time.to(self.device),
edge_index.to(self.device),
edge_type.to(self.device),
)
res = self.classifier.forward(node_rep[x_ids])
loss = self.criterion(res, ylabel.to(self.device))
self.optimizer.zero_grad()
torch.cuda.empty_cache()
loss.backward()
torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.clip)
self.optimizer.step()
train_losses += [loss.cpu().detach().tolist()]
self.train_step += 1
self.scheduler.step(self.train_step)
del res, loss
"""
Valid
"""
self.model.eval()
with torch.no_grad():
(
node_feature,
node_type,
edge_time,
edge_index,
edge_type,
x_ids,
ylabel,
) = valid_data
node_rep = self.gnn.forward(
node_feature.to(self.device),
node_type.to(self.device),
edge_time.to(self.device),
edge_index.to(self.device),
edge_type.to(self.device),
)
res = self.classifier.forward(node_rep[x_ids])
loss = self.criterion(res, ylabel.to(self.device))
"""
Calculate Valid F1. Update the best model based on highest F1 score.
"""
valid_f1 = f1_score(
res.argmax(dim=1).cpu().tolist(), ylabel.tolist(), average="micro"
)
if valid_f1 > self.best_val:
self.best_val = valid_f1
# torch.save(
# self.model,
# os.path.join(
# self.args.model_dir,
# self.args.task_name + "_" + self.args.conv_name,
# ),
# )
self.best_model_dict = deepcopy(self.model.state_dict())
print("UPDATE!!!")
self.st = time.time()
print(
(
"Epoch: %d (%.1fs) LR: %.5f Train Loss: %.2f Valid Loss: %.2f Valid F1: %.4f"
)
% (
epoch,
(self.st - self.et),
self.optimizer.param_groups[0]["lr"],
np.average(train_losses),
loss.cpu().detach().tolist(),
valid_f1,
)
)
self.stats += [[np.average(train_losses), loss.cpu().detach().tolist()]]
del res, loss
del train_data, valid_data
self.model.load_state_dict(self.best_model_dict)
best_model = self.model.to(self.device)
# best_model = torch.load(
# os.path.join(
# self.args.model_dir, self.args.task_name + "_" + self.args.conv_name
# )
# ).to(self.device)
best_model.eval()
gnn, classifier = best_model
with torch.no_grad():
test_res = []
for _ in range(10):
(
node_feature,
node_type,
edge_time,
edge_index,
edge_type,
x_ids,
ylabel,
) = node_classification_sample(
self.args,
self.target_type,
randint(),
self.test_target_nodes,
{1: True},
)
paper_rep = gnn.forward(
node_feature.to(self.device),
node_type.to(self.device),
edge_time.to(self.device),
edge_index.to(self.device),
edge_type.to(self.device),
)[x_ids]
res = classifier.forward(paper_rep)
test_f1 = accuracy_score(
res.argmax(dim=1).cpu().tolist(), ylabel.tolist()
)
test_res += [test_f1]
return np.average(test_res)