Predict with pytorch lightning when using BCEWithLogitsLoss for training - pytorch-lightning

Question

I'm trying to see how my trained model would predict a single instance of y and have of list of predicted and actual y.
It seems I'm missing a few steps and I'm not sure how to implement the predict_step, here is what I currently have:
mutag = ptgeom.datasets.TUDataset(root='.', name='MUTAG')
train_idx, test_idx = train_test_split(range(len(mutag)), stratify=[m.y[0].item() for m in mutag], test_size=0.25)
train_loader = ptgeom.loader.DataLoader(mutag[train_idx], batch_size=32, shuffle=True)
test_loader = ptgeom.loader.DataLoader(mutag[test_idx], batch_size=32)
class MUTAGClassifier(ptlight.LightningModule):
def __init__(self):
# The model is just GCNConv --> GCNConv --> graph pooling --> Dropout --> Linear
super().__init__()
self.gc1 = ptgeom.nn.GCNConv(7, 256)
self.gc2 = ptgeom.nn.GCNConv(256, 256)
self.linear = torch.nn.Linear(256, 1)
def forward(self, x, edge_index=None, batch=None, edge_weight=None):
# Note: "edge_weight" is not used for training, but only for the explainability part
if edge_index == None:
x, edge_index, batch = x.x, x.edge_index, x.batch
x = F.relu(self.gc1(x, edge_index, edge_weight))
x = F.relu(self.gc2(x, edge_index, edge_weight))
x = ptgeom.nn.global_mean_pool(x, batch)
x = F.dropout(x)
x = self.linear(x)
return x
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
return optimizer
def training_step(self, batch, _):
y_hat = self.forward(batch.x, batch.edge_index, batch.batch)
loss = F.binary_cross_entropy_with_logits(y_hat, batch.y.unsqueeze(1).float())
self.log("train_loss", loss)
self.log("train_accuracy", accuracy(y_hat, batch.y.unsqueeze(1)), prog_bar=True, batch_size=32)
return loss
def validation_step(self, batch, _):
x, edge_index, batch_idx = batch.x, batch.edge_index, batch.batch
y_hat = self.forward(x, edge_index, batch_idx)
self.log("val_accuracy", accuracy(y_hat, batch.y.unsqueeze(1)), prog_bar=True, batch_size=32)
checkpoint_callback = ptlight.callbacks.ModelCheckpoint(
dirpath='./checkpoints/',
filename='gnn-{epoch:02d}',
every_n_epochs=50,
save_top_k=-1)
trainer = ptlight.Trainer(max_epochs=200, callbacks=[checkpoint_callback])
trainer.fit(gnn, train_loader, test_loader)

Answer 1

The crux here is that you use F.binary_cross_entropy_with_logits in your training_step (for numerical stability I suppose). This means that nn.Sigmoid has to be applied to your output both in validation_step and predict_step as the operation is not part of forward(). Check this for more information. Notice that you may also need to round your predicted results depending on which accuracy method you are using in order to get correct metric results.
class MUTAGClassifier(ptlight.LightningModule):
def __init__(self):
# The model is just GCNConv --> GCNConv --> graph pooling --> Dropout --> Linear
super().__init__()
self.gc1 = ptgeom.nn.GCNConv(7, 256)
self.gc2 = ptgeom.nn.GCNConv(256, 256)
self.linear = torch.nn.Linear(256, 1)
self.s = nn.Sigmoid()
def forward(self, x, edge_index=None, batch=None, edge_weight=None):
# Note: "edge_weight" is not used for training, but only for the explainability part
if edge_index == None:
x, edge_index, batch = x.x, x.edge_index, x.batch
x = F.relu(self.gc1(x, edge_index, edge_weight))
x = F.relu(self.gc2(x, edge_index, edge_weight))
x = ptgeom.nn.global_mean_pool(x, batch)
x = F.dropout(x)
x = self.linear(x)
return x
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
return optimizer
def training_step(self, batch, _):
y_hat = self.forward(batch.x, batch.edge_index, batch.batch)
loss = F.binary_cross_entropy_with_logits(y_hat, batch.y.unsqueeze(1).float())
self.log("train_loss", loss)
self.log("train_accuracy", accuracy(y_hat, batch.y.unsqueeze(1)), prog_bar=True, batch_size=32)
return loss
def validation_step(self, batch, _):
x, edge_index, batch_idx = batch.x, batch.edge_index, batch.batch
y_hat = self.forward(x, edge_index, batch_idx)
y_hat = self.s(y_hat)
y_hat = torch.where(y_hat > 0.5, 1, 0) # may be needed
self.log("val_accuracy", accuracy(y_hat, batch.y.unsqueeze(1)), prog_bar=True, batch_size=32)
def predict_step(self, batch, _):
x, edge_index, batch_idx = batch.x, batch.edge_index, batch.batch
y_hat = self.forward(x, edge_index, batch_idx)
y_hat = self.s(y_hat)
y_hat = torch.where(y_hat > 0.5, 1, 0) # may be needed
return y_hat
You could then do the following in order to get a list of predictions with their corresponding ground truth:
batch = next(iter(train_loader)) # get a batch
y_hat = trainer.predict(your_model, batch)
print(list(y_hat))
print(list(batch.y))