Exporting PyTorch Lightning model to ONNX format not working

[pratikchhapolika]

I am using Jupyter Lab to run. It has pre-installed tf2.3_py3.6 kernel installed in it. It has 2 GPUS in it.

PyTorch Lightning Version (e.g., 1.3.0): '1.4.6'
PyTorch Version (e.g., 1.8): '1.6.0+cu101'
Python version: 3.6
OS (e.g., Linux): system='Linux'
CUDA/cuDNN version: 11.2
How you installed PyTorch (conda, pip, source): pip

Here is the screenshot of my model and it got interrupted due to connection issue.

I am saving the best model in checkpoint.

I am doing multi-label classification using Hugging face model. After training the model I want to export the model using ONNX format.

Here is the DataModule Class

N_EPOCHS = 30
BATCH_SIZE = 10

class  SRDataModule(pl.LightningDataModule):
    
    def __init__(self, X_train,y_train, X_test,y_test, tokenizer, batch_size=8, max_token_len=512):
        super().__init__()
        self.batch_size = batch_size
        self.train_df = X_train
        self.test_df = X_test
        self.train_lab = y_train
        self.test_lab = y_test
        self.tokenizer = tokenizer
        self.max_token_len = max_token_len

    def setup(self, stage=None):
        self.train_dataset = SRDataset(
          self.train_df,
          self.train_lab,
          self.tokenizer,
          self.max_token_len
        )

        self.test_dataset = SRDataset(
          self.test_df,
          self.test_lab,
          self.tokenizer,
          self.max_token_len
    )

    def train_dataloader(self):
        return DataLoader(
          self.train_dataset,
          batch_size=self.batch_size,
          shuffle=True,
          num_workers=10
        )

    def val_dataloader(self):
        return DataLoader(
          self.test_dataset,
          batch_size=self.batch_size,
          num_workers=10
        )

    def test_dataloader(self):
        return DataLoader(
          self.test_dataset,
          batch_size=self.batch_size,
          num_workers=10
        )

Here is the model class:

class SRTagger(pl.LightningModule):

  def __init__(self, n_classes: int, n_training_steps=None, n_warmup_steps=None):
    super().__init__()
    self.save_hyperparameters()
    self.bert = BertModel.from_pretrained(BERT_MODEL_NAME, return_dict=True)
    self.classifier = nn.Linear(self.bert.config.hidden_size, n_classes)
    self.n_training_steps = n_training_steps
    self.n_warmup_steps = n_warmup_steps
    self.criterion = nn.BCELoss()

  def forward(self, input_ids, attention_mask, labels=None):
    output = self.bert(input_ids, attention_mask=attention_mask)
    output = self.classifier(output.pooler_output)
    output = torch.sigmoid(output)    
    loss = 0
    if labels is not None:
        loss = self.criterion(output, labels)
    return loss, output

  def training_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    loss, outputs = self(input_ids, attention_mask, labels)
    
    self.log("train_loss", loss, prog_bar=True, logger=True)
    return {"loss": loss, "predictions": outputs, "labels": labels}

  def validation_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log("val_loss", loss, prog_bar=True, logger=True)
    return loss

  def test_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log("test_loss", loss, prog_bar=True, logger=True)
    return loss

  def training_epoch_end(self, outputs):
    
    labels = []
    predictions = []
    for output in outputs:
      for out_labels in output["labels"].detach().cpu():
        labels.append(out_labels)
      for out_predictions in output["predictions"].detach().cpu():
        predictions.append(out_predictions)

    labels = torch.stack(labels).int()
    predictions = torch.stack(predictions)

    for i, name in enumerate(LABEL_COLUMNS):
      class_roc_auc = auroc(predictions[:, i], labels[:, i])
      self.logger.experiment.add_scalar(f"{name}_roc_auc/Train", class_roc_auc, self.current_epoch)


  def configure_optimizers(self):

    optimizer = optim.RAdam(self.parameters(), lr=2e-4)

    scheduler = get_linear_schedule_with_warmup(
      optimizer,
      num_warmup_steps=self.n_warmup_steps,
      num_training_steps=self.n_training_steps
    )

    return dict(
      optimizer=optimizer,
      lr_scheduler=dict(
        scheduler=scheduler,
        interval='step'
      )
    )

Sample Data

sample_batch = next(iter(DataLoader(train_dataset, batch_size=10, num_workers=2)))
sample_batch["input_ids"].shape, sample_batch["attention_mask"].shape

(torch.Size([10, 512]), torch.Size([10, 512]))

sample_batch.keys()
dict_keys(['text_data', 'input_ids', 'attention_mask', 'labels'])

Model

model = SRTagger(
  n_classes=100,
  n_warmup_steps=warmup_steps,
  n_training_steps=total_training_steps 
)

ONNX code

# # Export the model
torch.onnx.export(model,                     # model being run
                  ##since model is in the cuda mode, input also need to be
                  (sample_batch["input_ids"],sample_batch["attention_mask"]),              # model input (or a tuple for multiple inputs)
                  "model_torch_export.onnx", # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=10,          # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names = ['input'],   # the model's input names
                  output_names = ['output'], # the model's output names
                  dynamic_axes={'input' : {0 : 'batch_size'},    # variable lenght axes
                                'output' : {0 : 'batch_size'}})

Error

RuntimeError: output 1 (0
[ CPULongType{} ]) of traced region did not have observable data dependence with trace inputs; this probably indicates your program cannot be understood by the tracer.

[awaelchli]

Hi

A google search reveals some help on this issue here:

pytorch/pytorch#31591

Citing the thread there

As the error message indicates, the tracer detected that the output of your model didn't have any relationship to the input.

If we look closer at your code, we see that loss=0 and labels=None.

  def forward(self, input_ids, attention_mask, labels=None):
    output = self.bert(input_ids, attention_mask=attention_mask)
    output = self.classifier(output.pooler_output)
    output = torch.sigmoid(output)    
    loss = 0
    if labels is not None:
        loss = self.criterion(output, labels)
    return loss, output

the if condition does not hold, so the part of your output (the loss) cannot be traced back to any inputs by onnx.

Change your code to something like this and try again please:

    if labels is not None:
        loss = self.criterion(output, labels)
        return loss, output
    return output

[pratikchhapolika]

Hi

A google search reveals some help on this issue here:

pytorch/pytorch#31591

Citing the thread there

As the error message indicates, the tracer detected that the output of your model didn't have any relationship to the input.

If we look closer at your code, we see that loss=0 and labels=None.

  def forward(self, input_ids, attention_mask, labels=None):
    output = self.bert(input_ids, attention_mask=attention_mask)
    output = self.classifier(output.pooler_output)
    output = torch.sigmoid(output)    
    loss = 0
    if labels is not None:
        loss = self.criterion(output, labels)
    return loss, output

the if condition does not hold, so the part of your output (the loss) cannot be traced back to any inputs by onnx.

Change your code to something like this and try again please:

    if labels is not None:
        loss = self.criterion(output, labels)
        return loss, output
    return output

@awaelchli Ok. This solves my problem. However I have one clarification :

In code to save in ONNX format:

torch.onnx.export(model,                     # model being run
                  ##since model is in the cuda mode, input also need to be
                  (sample_batch["input_ids"],sample_batch["attention_mask"]),              # model input (or a tuple for multiple inputs)
                  "model_torch_export.onnx", # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=10,          # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names = ['input'],   # the model's input names
                  output_names = ['output'], # the model's output names
                  dynamic_axes={'input' : {0 : 'batch_size'},    # variable lenght axes
                                'output' : {0 : 'batch_size'}})

Is this correct way to pass (sample_batch["input_ids"],sample_batch["attention_mask"]) inputs. Also what does

input_names = ['input'], # the model's input names output_names = ['output'], # the model's output names dynamic_axes={'input' : {0 : 'batch_size'}, # variable lenght axes 'output' : {0 : 'batch_size'}}) means?

Secondly, does this saves my best model which is trained above using pytorch lightning??

[awaelchli]

Great! Glad the main issue is resolved. I'm converting this issue to a discussion thread.

Also check out the PL docs for the to_onnx method on the LightningModule

Is this correct way to pass (sample_batch["input_ids"],sample_batch["attention_mask"]) inputs.

Yes, looks good.

Secondly, does this saves my best model which is trained above using pytorch lightning??

When you run trainer.fit(model) it will train but when it finishes you may not end up with the best weights. You can load the model with the best weights by accessing the path

trainer.checkpoint_callback.best_model_path

Check the ModelCheckpoint callback docs how to use it :)