Predict with pytorch lightning when using BCEWithLogitsLoss for training

user_10506634 0
Viewed 126

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)
pytorch-lightning
1 Answers

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))
user_95067730
Related