Common Pitfalls in Machine Learning Model Inference for Beginners and How to Solve Them
Tags: machinelearning, beginners, python, ai
Published: 2024-06-09 01:10:50

When training machine learning models and applying them to inference tasks, it's not uncommon to encounter issues, especially when switching computing environments (e.g., training on a GPU and inferring on a CPU). This often results in unstable predictions, such as alternating between two classes (0 and 1).

Here, I'll briefly summarize the issues I encountered and how I resolved them.

## 1. Model Loading Mistake
I built an inference endpoint on a CPU (`ml.m5.large`) and found that the model, which was trained on a GPU (`g4dn.2xlarge`), did not produce the expected inference results.

Upon checking the logs, I encountered the following warning:

```
2024-04-25T05:47:04,365 [WARN ] W-9000-model_1.0-stderr MODEL_LOG - Some weights of the model checkpoint at /opt/ml/model/code/pytorch_model.bin were not used when initializing BertModel...
```

This issue occurred when I attempted to directly load the `pytorch_model.bin` file, which was output from the fine-tuned `PredictionModel`, using `BertModel.from_pretrained`. The `BertModel.from_pretrained` method assumes the structure of the basic BERT model, thus it neglected the parameters of the LSTM and linear layers added to the `PredictionModel`, resulting in important parameters being overlooked.

```python
pretrained_config = path.join("/opt/ml/model/code/", "config.json")
pretrained_model = path.join("/opt/ml/model/code/", "pytorch_model.bin")
config = BertConfig.from_pretrained(pretrained_config)

model = PredictionModel(config=config, pretrained_model=pretrained_model)
```

The issue was resolved by properly loading the state from the `fine_tuning_model.pt`, which contained all the parameters of the model:

```python
model = PredictionModel(config=config, pretrained_model=None)
model_path = path.join("/opt/ml/model/code/", "fine_tuning_model.pt")
model.load_state_dict(torch.load(model_path))
```

## 2. Device Assignment for the Model
The model was initially set to use a CUDA device by default, which led to errors in an environment not supporting CUDA.

```
2024-04-28T06:59:31,905 [INFO ] W-9001-model_1.0-stdout MODEL_LOG - Exception in model fn Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False...
```

In an environment with only CPUs, it was necessary to appropriately assign the model to the correct device:

```python
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
```

## 3. Disabling Gradient Calculations
During inference, it is recommended to disable gradient calculations when model input data is provided as tensors. While I believe this had no effect on the inference results, failing to do so led to unnecessary memory use and increased computation time.

```python
with torch.no_grad():
    model_input = torch.tensor([preprocessed_data.getitem()], dtype=torch.long).to(device)
    model_output = model(model_input)
```

## Conclusion
Successfully transitioning from training to inference requires addressing key challenges like model loading, device compatibility, and efficient resource management. These solutions ensure more accurate and efficient machine learning applications.