Not able to Generate Predictions with Trainer.predict()

[spranjal25]

Hi, I'm new to PyTorch Lightning, used it for the first time and kind of liked it. However, I am facing this one problem, Implemented a classification task for which I trained the model with Huggingface pretrained model as base and classification head on top. The model is training successfully and giving decent validation losses. The problem is, I'm not quite able to figure out the inferencing part.

can anyone please point out what is it that I'm doing wrong? It's probably something very basic.

I'll add the classes of the lightning modules and the Data Modules below.

# |¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯|
# | Define the Pytorch Lightning Module Classifier Class     |
# |__________________________________________________________|

class ABSASentimentClassifier(pl.LightningModule):

  def __init__(self, learning_rate = setup['lr'], weights=None, **kwargs):
    super().__init__()

    self.save_hyperparameters('learning_rate', 'max_epochs')
    self.model = ABSAModel_Bert()
    self.weights = weights
    self.preds = []
  
  def training_step(self, batch, batch_nb):

    # Forward
    y_hat = self.model(batch)

    # if self.weights:
    #   self.weights = torch.tensor(class_weights,dtype=torch.float) 
    
    # Loss
    loss_fct = torch.nn.CrossEntropyLoss()
    
    loss = loss_fct(y_hat.view(-1, self.model.num_labels), batch['label'].view(-1))

    # Logs
    self.log_dict({'training_loss':loss}, prog_bar=True)

    return loss

  
  def validation_step(self, batch, batch_nb):
    
    # Forward
    y_hat = self.model(batch)
        
    # Loss
    loss_fct = torch.nn.CrossEntropyLoss()
    loss = loss_fct(y_hat.view(-1, self.model.num_labels), batch['label'].view(-1))

    # Acc
    a, y_hat = torch.max(y_hat, dim=1)
    val_acc = accuracy_score(y_hat.cpu(), batch['label'].cpu())
    val_acc = torch.tensor(val_acc)
    
    # Logs
    self.log_dict({'val_loss':loss,'val_acc':val_acc}, prog_bar=True)
    
    return loss

  
  def test_step(self, batch, batch_nb):
    self.model.eval()
    
    # Forward
    yhat = self.model(batch)
      
    # Loss
    # loss_fct = torch.nn.CrossEntropyLoss()
    # loss = loss_fct(y_hat.view(-1, self.model.num_labels), batch['label'].view(-1))
    
    # a, y_hat = torch.max(y_hat, dim=1)
    # test_acc = accuracy_score(y_hat.cpu(), batch['label'].cpu())
    
    # Logs
    # self.log_dict({'test_loss':loss,'test_acc':test_acc}, prog_bar=True)
    self.preds = self.preds.extend(yhat.cpu().detach().numpy().tolist())
    return 

  
  def predict_dataloader(self, batch, batch_idx: int , dataloader_idx: int = None):

    return self.model(batch)


  '''
  |¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯|
  | Training Setup  |
  |_________________|
  '''
  def configure_optimizers(self):
    '''
    |¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯|
    |   REQUIRED                                                            |
    |   can return multiple optimizers and learning_rate schedulers         |
    |   (LBFGS it is automatically supported, no need for closure function) |
    |_______________________________________________________________________|
    '''
    optimizer = torch.optim.Adam([p for p in self.parameters() if p.requires_grad], lr=self.hparams.learning_rate, eps=1e-08)
    scheduler = {   
      'scheduler': torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=2e-5, 
                                                       steps_per_epoch=len(self.trainer.datamodule.train_dataloader()),
                                                       epochs=self.hparams.max_epochs),
                 
      'interval': 'step'  # called after each training step
    } 
    
    return [optimizer], [scheduler]

  @staticmethod
  def add_model_specific_args(parent_parser, root_dir):  # pragma: no-cover
    """
    |¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯|
    | Define parameters that only apply to this model     |
    |_____________________________________________________|
    """
    parser = ArgumentParser(parents=[parent_parser])

    # data
    parser.add_argument('--data_root', default=os.path.join(root_dir, 'train_val_data'), type=str)

    # training params (opt)
    parser.add_argument('--learning_rate', default=setup['lr'], type=float, help = "type (default: %(default)f)")
    return parser

also the dataset class is :

class ABSADataset(Dataset):
  def __init__(self, df, tokenizer, max_len=setup['max_sen_length']):
    self.texts = df['text']
    self.aspects = df['aspect']
    if 'label' in df.columns:
      # print('****Labels Present****')
      self.targets = df['label']

    else:
      self.targets = None

    self.tokenizer = tokenizer
    self.max_len = max_len

  def __len__(self):
    return len(self.aspects)

  def __getitem__(self, idx):

    # convert indexes, tensor->list
    if torch.is_tensor(idx):
      idx = idx.tolist()
    
    # define the aspect and text item
    text = (str(self.texts[idx]))
    aspect = str(self.aspects[idx])

    # define the label
    target = self.targets[idx]

    # pair the aspect and text for pair-encoding
    pairs = [text, aspect]
    
    '''
    # |¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯|
    # | For Debugging            |
    # |__________________________|
    # print(f' text: {text}')
    # print(f' aspect: {aspect}')
    # print(type(text))
    # print(type(aspect))
    '''
    
    # encode the feature pair
    encoded = self.tokenizer.encode_plus(pairs,
                                    add_special_tokens=True,
                                    padding='max_length', 
                                    max_length=setup['max_sen_length'], 
                                    return_attention_mask=True,
                                    return_tensors='pt',
                                    truncation=True)
    '''
    # |¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯|
    # | For Debugging            |
    # |__________________________|
    # for ids in encoded['input_ids']:
    #   print('*'*20)
    #   print(f'{self.tokenizer.decode(ids)} of length = {len(self.tokenizer.decode(ids).split(" "))}')
    #   print(f'is encoded as : \n{ids} \nwith length = {len(ids)}')
    #   print('*'*20)
    '''
    
    return {
        'label' : target,
        'input_ids' : encoded['input_ids'],
        'attention_mask' : encoded['attention_mask'] 
    }

My goal is to be able to generate predictions for data without any labels present, using the trained model (saved as checkpoint (.ckpt))

This is what I did:

testset = ABSATest_Dataset(test, tokenizer=transformer_tokenizer)

testLoader = DataLoader(testset, batch_size=setup['test_batch_size'])

trainer.predict(model_infer, testLoader)

Where model_infer is :

model_infer = ABSASentimentClassifier.load_from_checkpoint(PATH_TO_CKPT_FILE)

and got :

---------------------------------------------------------------------------
MisconfigurationException                 Traceback (most recent call last)
<ipython-input-44-c724efd019b7> in <module>()
----> 1 trainer.predict(model_infer, testLoader)

5 frames
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in predict(self, model, dataloaders, datamodule, return_predictions, ckpt_path)
    987         """
    988         return self._call_and_handle_interrupt(
--> 989             self._predict_impl, model, dataloaders, datamodule, return_predictions, ckpt_path
    990         )
    991 

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in _call_and_handle_interrupt(self, trainer_fn, *args, **kwargs)
    680         """
    681         try:
--> 682             return trainer_fn(*args, **kwargs)
    683         # TODO: treat KeyboardInterrupt as BaseException (delete the code below) in v1.7
    684         except KeyboardInterrupt as exception:

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in _predict_impl(self, model, dataloaders, datamodule, return_predictions, ckpt_path)
   1030         )
   1031 
-> 1032         results = self._run(model, ckpt_path=self.predicted_ckpt_path)
   1033 
   1034         assert self.state.stopped

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in _run(self, model, ckpt_path)
   1115             parsing.clean_namespace(model.hparams)
   1116 
-> 1117         verify_loop_configurations(self, model)
   1118 
   1119         # attach model log function to callback

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/configuration_validator.py in verify_loop_configurations(trainer, model)
     38         __verify_eval_loop_configuration(trainer, model, "test")
     39     elif trainer.state.fn == TrainerFn.PREDICTING:
---> 40         __verify_eval_loop_configuration(trainer, model, "predict")
     41 
     42     __verify_dp_batch_transfer_support(trainer, model)

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/configuration_validator.py in __verify_eval_loop_configuration(trainer, model, stage)
    187             raise MisconfigurationException("`predict_step` cannot be None to run `Trainer.predict`")
    188         elif not has_step and not is_overridden("forward", model):
--> 189             raise MisconfigurationException("`Trainer.predict` requires `forward` method to run.")
    190     else:
    191         # -----------------------------------

MisconfigurationException: `Trainer.predict` requires `forward` method to run.

ALso, I haven't defined a forward function in the lightning module because it is present in the model class:

class ABSAModel_Bert(torch.nn.Module):

  def __init__(self, num_labels=setup['num_labels'], config = setup, **kwargs):
    super(ABSAModel_Bert, self).__init__()
    
    self.num_labels = num_labels
    self.bert = transformers.AutoModel.from_pretrained(config['model_name'])
    self.bert_config = transformers.AutoConfig.from_pretrained(config['model_name'])

    self.pre_classifier = torch.nn.Linear(self.bert_config.hidden_size, self.bert_config.hidden_size)

    self.classifier = torch.nn.Linear(self.bert_config.hidden_size, self.num_labels)

    self.dropout = torch.nn.Dropout(self.bert_config.hidden_dropout_prob)
    # print(f'Using Dropout = {self.bert.config.seq_classif_dropout}')

    self.relu = torch.nn.ReLU()

    '''
    |¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯|
    | freeze the layers of Bert for training if needed so that |   
    | the embeddings of all layers of Bert are not changed     |
    |__________________________________________________________|
    '''
    # for param in self.bert.parameters():
    #   param.requires_grad = False

  
  def forward(self, batch):

  #   print((batch['input_ids'].squeeze(1)).shape)
  #   print("*"*10)
  #   print(batch['input_ids'])
  #   print("*"*10)
    
    outputs = self.bert(input_ids=batch['input_ids'].squeeze(1), 
                        attention_mask=batch['attention_mask'])
    
    # output from last hidden layer
    hidden_state = outputs[0]  # (batch_size, seq_len, dim)

    '''
    |¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯|
    | *output of [CLS] token                                   |
    |                                                          |
    | [CLS] token contains the pooled embeddings of the entire | 
    | Sequence, these are used for the classification.         |
    |__________________________________________________________|
    '''
    pooled_output = hidden_state[:, 0] # (batch_size, dim)

    '''
    |¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯|
    | sending the [CLS] token embeddings through Linear, ReLU  |
    | and Dropout layers                                       |
    |__________________________________________________________|
    '''
    pooled_output = self.pre_classifier(pooled_output)  # (batch_size, dim)
    pooled_output = self.relu(pooled_output)  # (batch_size, dim)
    pooled_output = self.dropout(pooled_output)  # (batch_size, dim)
    logits = self.classifier(pooled_output)  # (batch_size, num_labels)

    return logits

  def get_outputs(self, input_ids, attention_mask):
    outputs = self.bert(input_ids=input_ids, \
                        attention_mask=attention_mask)
    

[rohitgr7]

since your model is an instance of your LightningModule it cannot rely on model.forward to generate the predictions because predict_step by default calls LightningModule.predict.
you need to either override predict_step

def predict_step(...):
    return self.model(...)

or define forward method in your lightning module

def forward(...):
    return self.model(...)

[spranjal25]

Hi @rohitgr7, thanks for the #LightningFast reply. I got it working by defining a predict_step() function. Could you also tell me that which method would be better for this? the predict_step() or defining a forward() function? Also if I define forward() method how do I generate predictions then?

model.eval()
preds = model(x)

# code to convert tensors to labels

is this the right way?

[rohitgr7]

Lightning calls forward by default inside predict_step so either is fine.

Also if I define forward() method how do I generate predictions then?

firs thing, you don't need to call model.eval() explicitly. And yes, this looks good. All you need to do is return the predictions from predict_step and it will get accumulated itself.