Getting positive prediction list

[anthonysirico]

Is there a way to get the positive predictions separated from the prediction list?

def predict(loader):
    model.eval()

    predictions = []
    pred_proba = []
    for data in loader:
        data = data.to(device)
        out = model(data.x.to(torch.float), data.edge_index, data.batch)
        pred_proba.append(torch.Tensor.tolist(F.softmax(out)[0]))
        pred = out.argmax(dim=1)
        predictions.append(pred.item())
    return predictions, pred_proba

In the code, I can create an index list to append to, but I was wondering if there was an easier way to set the positively predicted graph instance aside for later use.

[wsad1]

If predictions is a list I think saving the indicies is the best approach. If predicitons was a tensor you could simply do predictions[predictions>0].

[anthonysirico]

@wsad1 Thanks! one small thing, so now that I have a list of Data objects that were positive predictions, how can I convert a list to a Dataset?

[anthonysirico]

Elevating to the man himself, @rusty1s

[rusty1s]

Is there a need to convert your list of data sets to a Dataset? The DataLoader just works fine on a list of data objects. If you still want to create a Dataset, you can do:

class MyDataset(InMemoryDataset):
    def __init__(self, data_list: List[Data]):
        super().__init__(None)
        self.data, self.slices = self.collate(data_list)

[anthonysirico]

I was wondering because prior to uploading it to a dataloader, I wanted to access the data like you would with InMemory