UnpicklingError Weights only load failed stops model training

[john2408]

Context:
Python version 3.10.15 (using python venv environment)
Hardware: Chip: Apple M3 Max

Problem:
I am trying a simple Tabular model with FTTransformer, However it fails after training is finished. I already tried adding torch.serialization.add_safe_globals([DictConfig]), however the error still persists, even though the model is trained it fails when reloading the weights. I also tried with a higher python version (3.12), however it did not work.

Error:

Exception has occurred: UnpicklingError Weights only load failed. This file can still be loaded, to do so you have two options, �[1mdo those steps only if you trust the source of the checkpoint�[0m. (1) In PyTorch 2.6, we changed the default value of the weights_only argument in torch.load from False to True. Re-running torch.load with weights_only set to False will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source. (2) Alternatively, to load with weights_only=True please check the recommended steps in the following error message. WeightsUnpickler error: Unsupported global: GLOBAL omegaconf.base.ContainerMetadata was not an allowed global by default. Please use torch.serialization.add_safe_globals([ContainerMetadata]) or the torch.serialization.safe_globals([ContainerMetadata]) context manager to allowlist this global if you trust this class/function.

Check the documentation of torch.load to learn more about types accepted by default with weights_only [https://pytorch.org/docs/stable/generated/torch.load.html](vscode-file://vscode-app/Applications/Visual%20Studio%20Code.app/Contents/Resources/app/out/vs/code/electron-sandbox/workbench/workbench.html). _pickle.UnpicklingError: Unsupported global: GLOBAL omegaconf.base.ContainerMetadata was not an allowed global by default. Please use torch.serialization.add_safe_globals([ContainerMetadata]) or the torch.serialization.safe_globals([ContainerMetadata]) context manager to allowlist this global if you trust this class/function.

Code:

  
        # Convert target column to int8
    df_train[target_col] = df_train[target_col].astype('int32')
    df_test[target_col] = df_test[target_col].astype('int32')
    
    # Split data into train and test
    train, val = train_test_split(df_train, test_size=0.1, random_state=random_seed)

    x_test= df_test[features]
    y_test = df_test[target_col]

    # Config 
    data_config = DataConfig(
        target=target_col,  # target should always be a list.
        continuous_cols=num_cols,
        categorical_cols=cat_cols,
    )

    trainer_config = TrainerConfig(
        #     auto_lr_find=True, # Runs the LRFinder to automatically derive a learning rate
        batch_size=32,
        max_epochs=30,
        early_stopping="valid_loss",  # Monitor valid_loss for early stopping
        early_stopping_mode="min",  # Set the mode as min because for val_loss, lower is better
        early_stopping_patience=5,  # No. of epochs of degradation training will wait before terminating
        checkpoints="valid_loss",  # Save best checkpoint monitoring val_loss
        load_best=True,  # After training, load the best checkpoint
    )

    optimizer_config = OptimizerConfig()

    head_config = LinearHeadConfig(
        layers="",  # No additional layer in head, just a mapping layer to output_dim
        dropout=0.1,
        initialization="kaiming",
    ).__dict__ 


    # FT Tabular Transformer
    model_config = FTTransformerConfig(
        task="classification",
        learning_rate=1e-3,
        head="LinearHead",  # Linear Head
        head_config=head_config,  # Linear Head Config
    )

    tabular_model = TabularModel(
        data_config=data_config,
        model_config=model_config,
        optimizer_config=optimizer_config,
        trainer_config=trainer_config,
    )
    
        
    tabular_model.fit(train=train, validation=val)
    

[dosubot[bot]]

Hey @john2408! I'm here to help you with your coding issues. I can assist in solving bugs, answering questions, and becoming a contributor. Let's tackle this problem together!

The error you're encountering is related to the changes in PyTorch 2.6, where the default value of the weights_only argument in torch.load was changed from False to True. This change affects how models are loaded, particularly when certain globals are not allowed by default.

To resolve this issue, you can try the following steps:

1. Add ContainerMetadata to Safe Globals: The error message suggests adding ContainerMetadata to the safe globals. You can do this by using torch.serialization.add_safe_globals([ContainerMetadata]) or the torch.serialization.safe_globals([ContainerMetadata]) context manager. This will allow the ContainerMetadata global to be used during the unpickling process.

2. Set weights_only=False: If you trust the source of the checkpoint, you can set weights_only=False when calling torch.load. This will revert to the previous behavior and might resolve the issue. However, be cautious as this can result in arbitrary code execution if the checkpoint is not from a trusted source.

Here's how you can implement these changes:

import torch
from omegaconf.base import ContainerMetadata

# Add ContainerMetadata to safe globals
torch.serialization.add_safe_globals([ContainerMetadata])

# When loading the model, set weights_only=False if you trust the source
model = torch.load('path_to_model_checkpoint', weights_only=False)

Make sure to replace 'path_to_model_checkpoint' with the actual path to your model checkpoint file.

These steps should help you load the model without encountering the UnpicklingError.

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Other}

[john2408]

Thanks for the reply, I had to implement the following changes to make it work:

In my training file:

from omegaconf import DictConfig
torch.serialization.safe_globals([DictConfig])

in the file pytorch_tabular/utils/python_utils.py adding the parameters weights_only=False:

def pl_load(
    path_or_url: Union[IO, _PATH],
    map_location: _MAP_LOCATION_TYPE = None,
) -> Any:
    """Loads a checkpoint.

    Args:
        path_or_url: Path or URL of the checkpoint.
        map_location: a function, ``torch.device``, string or a dict specifying how to remap storage locations.

    """
    if not isinstance(path_or_url, (str, Path)):
        # any sort of BytesIO or similar
        return torch.load(path_or_url, map_location=map_location, weights_only=False)
    if str(path_or_url).startswith("http"):
        return torch.hub.load_state_dict_from_url(
            str(path_or_url),
            map_location=map_location,  # type: ignore[arg-type] # upstream annotation is not correct
        )
    fs = get_filesystem(path_or_url)
    with fs.open(path_or_url, "rb") as f:
        return torch.load(f, map_location=map_location, weights_only=False)

Additionally, I tried testing my code inside a Docker container with the bullseye image for python version 3.12.7 (LInux Debian). However it still did not work. It showed the same error until I did the above changes.

[ArozHada]

Hello @john2408 ,
This is an unpickling error that traces to the update with torch.load function in Pytorch 2.6 . You can either downgrade to a older Pytorch version or wait until PR #543 is merged.

[john2408]

Hi @ArozHada Thanks, I did the manual fix on my local setting weights_only=False in order to make it work. Thanks that the PR was raised.