Make a single prediction with pytorch geometric GCNN

1

Hello, I am a beginner with machine learning so please forgive me if this is a stupid question.

I’m trying to use a graph convolutional neural network to predict the classification of 3D data, specifically cell morphology. This is my testing method, where target is a one dimensional matrix of size n, n being the number of vertices.

def test(model, test_loader, num_nodes, target, device):
model.eval()
correct = 0
total_loss = 0
n_graphs = 0
with torch.no_grad():
for idx, data in enumerate(test_loader):
out = model(data.to(device))
total_loss += F.nll_loss(out, target).item()
pred = out.max(1)[1]
correct += pred.eq(target).sum().item()
n_graphs += data.num_graphs
return correct / (n_graphs * num_nodes), total_loss / len(test_loader)

How could I produce a single prediction for a piece of data instead of the tensor of predictions?

2

Assuming your input uses a shape of [batch_size, *], you could set the batch_size to 1 and pass this single sample to the model.
The rest of the code should stay the same, as the used method should not depend on the actual batch size.