#Print model summary
print(model)
map = ('cpu' if device == 'cpu' else None)
# Pass graphs to model and get node embeddings
model.eval()
counter = 0
with torch.no_grad():
for batch_idx, (graph, pdb_ids) in enumerate(loader):
print('batch n° ', batch_idx, '/', len(loader))
graph = send_graph_to_device(graph, device)
# Forward pass
model.GNN(graph)
#Unbatch and get node embeddings
graphs = dgl.unbatch(graph)
for i, g in enumerate(graphs):
gid = pdb_ids[i]
# Load networkx original graph
with open(f'{args.train_dir}/{gid}.pickle', 'rb') as f:
fr3d_g = pickle.load(f)
node_map = {
nid: old_id
for nid, old_id in enumerate(sorted(fr3d_g.nodes()))
}
for epoch in range(1, args.epochs + 1):
print('*********************************')
print(f'Starting epoch {epoch}')
train_ep_loss, test_ep_loss = 0, 0
pos_pred = 0
for batch_idx, (graph, pdbids, labels) in enumerate(train_loader):
total_steps += 1 # count training steps
graph = send_graph_to_device(graph, device)
labels = torch.tensor(labels).to(device)
# Warm start embeddings
if (args.embeddings):
init_embeddings.GNN(graph)
# Forward pass
h = model(graph)
#Compute loss
t_loss = criterion(m(h), labels)
optimizer.zero_grad()
t_loss.backward()
# Epoch accuracy
_, pred = torch.max(m(h), dim=1)
pos_pred += torch.sum(pred).float()
# Confusion matrix :
true, pred = labels.cpu().detach(), pred.cpu().detach()