Fine Tuning error when fetching model using load_model

[Draghont]

Describe the bug
I'm am tring to train a model using SSL, pretrain and finetune as described here :
https://pytorch-tabular.readthedocs.io/en/latest/tutorials/08-Self-Supervised%20Learning-DAE/

My pretrain and finetune function works well in case the pretrain pass the ssl model directly to finetune.
The script return the following error in case I save the SSL model and then load. ( this is usefull to split train time and avoid having two long train run for each phase).

Traceback (most recent call last):
File "pytorch_tabular_example.py", line 725, in
finetuned_model = run_fine_tuning_v1(ssl_model, train_data, val_data, target_column, run_finetuning=True)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "pytorch_tabular_example.py", line 662, in run_fine_tuning_v1
finetune_model = ssl_model.create_finetune_model(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "Lib\site-packages\pytorch_tabular\tabular_model.py", line 1057, in create_finetune_model
model = model_callable(
^^^^^^^^^^^^^^^
File "Lib\site-packages\pytorch_tabular\models\base_model.py", line 676, in init
super().init(
File "\Lib\site-packages\pytorch_tabular\models\base_model.py", line 169, in init
self._build_network()
File "\Lib\site-packages\pytorch_tabular\models\base_model.py", line 718, in _build_network
self._head = self._get_head_from_config()
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "MLTrn\Lib\site-packages\pytorch_tabular\models\base_model.py", line 696, in _get_head_from_config
return _head_callable(
^^^^^^^^^^^^^^^
File "Lib\site-packages\pytorch_tabular\models\common\heads\blocks.py", line 60, in init
_layers.append(nn.Linear(_curr_units, output_dim))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "Lib\site-packages\torch\nn\modules\linear.py", line 99, in init
self.weight = Parameter(torch.empty((out_features, in_features), **factory_kwargs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: empty() received an invalid combination of arguments - got (tuple, dtype=NoneType, device=NoneType), but expected one of:

• (tuple of ints size, *, tuple of names names, torch.memory_format memory_format = None, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)
• (tuple of ints size, *, torch.memory_format memory_format = None, Tensor out = None, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)

To Reproduce
pretrain the model
save using

def save_model(model, path):
    model.save_model(path)
    log_step("Save Model", f"Model saved at {path}")

load it using:

def load_model_v2(path, device='cuda'):
map_location = torch.device(device) if torch.cuda.is_available() else torch.device('cpu')
model = TabularModel.load_model(path, map_location=map_location)
log_step("Load Model", f"Model loaded from {path} on device {map_location}")
return model

and pass it to the same finetuing function that when provided with a model not saved works.

Expected behavior
I can't understand why saving and loading the model and the fine tuining it cause an error while passing the SSL model directly do not cause any issue. Having the ability to save and load the model make debugging and retrain much quicker.

Regards

[dosubot]

Hey @Draghont! I'm here to help you with any bugs, questions, or contributions you have. Let's tackle this issue together!

The error you're encountering seems to be related to the model's state or configuration not being correctly preserved or restored when saving and loading the model. This discrepancy can occur if certain attributes or configurations are not saved or if they are not correctly reinitialized upon loading.

Here are a few things you can check or try:

1. Ensure Complete State Saving: Verify that the save_model method in the TabularModel class saves all necessary components of the model, including the model architecture, weights, and any additional configurations or hyperparameters. If any part of the model's state is not saved, it could lead to issues when loading.

2. Check Device Mapping: When loading the model, ensure that the map_location parameter is correctly mapping the model to the intended device. This is crucial if the model was trained on a different device than the one it is being loaded onto ^[1].

3. Verify Head Configuration: The error seems to originate from the _get_head_from_config method, which is responsible for setting up the model's head based on the configuration. Ensure that the head configuration is correctly specified and that all necessary parameters are provided when creating the fine-tune model ^[2].

4. Consistency in Model Initialization: When you load the model, ensure that any initialization logic that was applied when the model was first created is also applied after loading. This includes setting up the network architecture and any custom layers or components.

5. Debugging the Head Initialization: The specific error message indicates an issue with the dimensions provided to torch.empty(). Check the dimensions being passed to the nn.Linear layer in the head configuration to ensure they are valid and consistent with the model's architecture.

By addressing these areas, you should be able to identify the root cause of the issue and ensure that the model can be saved, loaded, and fine-tuned without errors.

To continue talking to Dosu, mention @dosu.

[Draghont]

train_data, val_data, dataset_params, ssl_train, ssl_val = load_and_create_datasets(db_manager_ML_analysis, timeseries=False, return_type='dataframe', return_full_dataset = True)

batch_size = 2048
steps_per_epoch = int(ssl_train.shape[0]/batch_size)
epochs = 10

checkpoint_dir = "best_model_save_load_test"  # Directory to save/load best model checkpoint
SSL_checkpoint_path = os.path.join(checkpoint_dir, "SSL_model.pth")
fine_tuned_path = os.path.join(checkpoint_dir, "fine_tuned_model.pth")

target_column = 'score'
continuous_columns = [col for col in train_data.columns if col != target_column]

def log_step(step, message):
    print(f"[{step}] {message}")

# Save model function
def save_model(model, path):
    model.save_model(path)
    log_step("Save Model", f"Model saved at {path}")
    
# Load model function with proper GPU mapping
def load_model_v2(path, device='cuda'):
    # Ensure device compatibility
    map_location = torch.device(device) if torch.cuda.is_available() else torch.device('cpu')
    
    # Load model with specified map location
    model = TabularModel.load_model(path, map_location=map_location)
    
    # Log and return the model
    log_step("Load Model", f"Model loaded from {path} on device {map_location}")
    return model
SSL = True
if SSL:
    ssl_data_config = DataConfig(
        target=None, #Setting target as None because we don't need the target for SSL
        continuous_cols=continuous_columns,
        normalize_continuous_features=True,
        handle_missing_values=False, # For SSL tasks, missing values and unknwon categories will not be handled automatically
        handle_unknown_categories=False, # Not Setting these configs to False will throw and error when initializing TabularModel 
    )

    ssl_trainer_config = TrainerConfig(
        batch_size=batch_size,
        max_epochs=epochs,
        early_stopping="valid_loss", # Turning off Early Stopping
        checkpoints="valid_loss", # Save best checkpoint monitoring val_loss
        load_best=True, # After training, load the best checkpoint
    )

    # Setting OneCycleLR schedule
    ssl_optimizer_config = OptimizerConfig(
        lr_scheduler="OneCycleLR",
        lr_scheduler_params={
            "max_lr":1e-2, 
            "epochs": epochs, 
            "steps_per_epoch":steps_per_epoch
        }
    )

    # Setting the encoder config
    encoder_config = CategoryEmbeddingModelConfig(
        task="backbone",
        layers="4000-2000-1000-512",  # Number of nodes in each layer
        activation="ReLU",  # Activation between each layers
        head=None, # If not set to None, it will throw a warning
    )

    # Setting the decoder config.
    # NOTE: the last dimension in encoder layers should be first dimension in decoder layers
    # i.e. last encoder layer dim = 512, first decoder layer dim = 512
    decoder_config = CategoryEmbeddingModelConfig(
        task="backbone",
        layers="512-2048-4096",  # Number of nodes in each layer
        activation="ReLU",  # Activation between each layers
        head=None, # If not set to None, it will throw a warning
    )

    # DAE Config. No need to set task because it is hardcoded to SSL
    # Can't set any loss or metrics as well because for the SSL task
    # (especially for DAE), the loss and metrics are fixed.
    ssl_model_config = DenoisingAutoEncoderConfig(
        noise_strategy="swap", # Can be 'zero' as well. Defines if noise is swapping features or making features zero
        default_noise_probability = 0.7, # Probability of corruption by noise
        include_input_features_inference=True, # Setting this to True will include the input features in the concatenated layers in the decoder
        encoder_config=encoder_config, 
        decoder_config=decoder_config, 
        learning_rate=1e-3)

    ssl_tabular_model = TabularModel(
        data_config=ssl_data_config,
        model_config=ssl_model_config,
        optimizer_config=ssl_optimizer_config,
        trainer_config=ssl_trainer_config,
        verbose=False
    )

    # Pretraining
    ssl_tabular_model.pretrain(train=ssl_train, validation=ssl_val)
    save_model(ssl_tabular_model, SSL_checkpoint_path)

if ssl_tabular_model:
    ssl_tabular_model_loaded = ssl_tabular_model
else:
    ssl_tabular_model_loaded = load_model_v2(SSL_checkpoint_path)


ft_trainer_config = TrainerConfig(
    batch_size=2048,
    max_epochs=epochs,
    early_stopping="valid_loss", # Turning off Early Stopping
    checkpoints="valid_loss", # Save best checkpoint monitoring val_loss
    load_best=True, # After training, load the best checkpoint
)

ft_optimizer_config = OptimizerConfig(
    lr_scheduler="OneCycleLR",
    lr_scheduler_params={
        "max_lr":1e-3, 
        "epochs": epochs, 
        "steps_per_epoch":steps_per_epoch
    }
)

finetune_model = ssl_tabular_model_loaded.create_finetune_model(
    task="regression",
    train=train_data, 
    validation= val_data,  
    target=[target_column], #Provide the column name of the target as a list
    head="LinearHead",
    head_config={
        "layers": "256-64",
        "activation": "ReLU",
    },
    trainer_config=ft_trainer_config,# Overriding previous trainer config
    optimizer_config=ft_optimizer_config,
    optimizer="Adam", # Using a custom optimizer
    #optimizer_params={"nus": (0.7, 1.0), "betas": (0.95, 0.998)}
)

finetune_model.finetune(freeze_backbone=True)

This script is taken from documentation ( https://pytorch-tabular.readthedocs.io/en/latest/tutorials/08-Self-Supervised%20Learning-DAE/) I just added a step to save ssl_tabular_model and load with a different name, and this cause the problem described above.

    # Pretraining
    ssl_tabular_model.pretrain(train=ssl_train, validation=ssl_val)
    save_model(ssl_tabular_model, SSL_checkpoint_path)

if ssl_tabular_model:
    ssl_tabular_model_loaded = ssl_tabular_model
else:
    ssl_tabular_model_loaded = load_model_v2(SSL_checkpoint_path)

ERROR:

TypeError: empty() received an invalid combination of arguments - got (tuple, dtype=NoneType, device=NoneType), but expected one of:

• (tuple of ints size, *, tuple of names names, torch.memory_format memory_format = None, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)
• (tuple of ints size, *, torch.memory_format memory_format = None, Tensor out = None, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)

set SSL = True to complete training without error

[manujosephv]

@Draghont Looks like some issue with saving and restoring. I'll try to take a look at it this weekend

[manujosephv]

I tried with the covertype dataset and it doesn't throw and error:

from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score
import random
import numpy as np
import pandas as pd
from pathlib import Path
from pytorch_tabular.utils import print_metrics, load_covertype_dataset

data, cat_col_names, num_col_names, target_name = load_covertype_dataset()
# Splitting data into data for SSL and data for finetuning
ssl, finetune = train_test_split(data, random_state=42, test_size=0.001)
# Train and val splits
ssl_train, ssl_val = train_test_split(ssl, random_state=42)
finetune_train, finetune_test = train_test_split(finetune, random_state=42)
finetune_train, finetune_val = train_test_split(finetune_train, random_state=42)
print(f"Unlabelled Data: {ssl.shape[0]} rows | Labelled Data: {finetune.shape[0]}")

from pytorch_tabular import TabularModel
from pytorch_tabular.models import CategoryEmbeddingModelConfig
from pytorch_tabular.config import (
    DataConfig,
    OptimizerConfig,
    TrainerConfig,
    ExperimentConfig,
)
from pytorch_tabular.models.common.heads import LinearHeadConfig
from pytorch_tabular.ssl_models.dae import DenoisingAutoEncoderConfig


batch_size = 2048
steps_per_epoch = int(ssl_train.shape[0] / batch_size)
epochs = 100

ssl_data_config = DataConfig(
    target=None,  # Setting target as None because we don't need the target for SSL
    continuous_cols=num_col_names,
    categorical_cols=cat_col_names,
    normalize_continuous_features=True,
    handle_missing_values=False,  # For SSL tasks, missing values and unknwon categories will not be handled automatically
    handle_unknown_categories=False,  # Not Setting these configs to False will throw and error when initializing TabularModel
)

ssl_trainer_config = TrainerConfig(
    batch_size=batch_size,
    max_epochs=epochs,
    early_stopping="valid_loss",  # Turning off Early Stopping
    checkpoints="valid_loss",  # Save best checkpoint monitoring val_loss
    load_best=True,  # After training, load the best checkpoint
    fast_dev_run=True,
)

# Setting OneCycleLR schedule
ssl_optimizer_config = OptimizerConfig(
    lr_scheduler="OneCycleLR",
    lr_scheduler_params={
        "max_lr": 1e-2,
        "epochs": epochs,
        "steps_per_epoch": steps_per_epoch,
    },
)

# Setting the encoder config
encoder_config = CategoryEmbeddingModelConfig(
    task="backbone",
    layers="4000-2000-1000-512",  # Number of nodes in each layer
    activation="ReLU",  # Activation between each layers
    head=None,  # If not set to None, it will throw a warning
)

# Setting the decoder config.
# NOTE: the last dimension in encoder layers should be first dimension in decoder layers
# i.e. last encoder layer dim = 512, first decoder layer dim = 512
decoder_config = CategoryEmbeddingModelConfig(
    task="backbone",
    layers="512-2048-4096",  # Number of nodes in each layer
    activation="ReLU",  # Activation between each layers
    head=None,  # If not set to None, it will throw a warning
)

# DAE Config. No need to set task because it is hardcoded to SSL
# Can't set any loss or metrics as well because for the SSL task
# (especially for DAE), the loss and metrics are fixed.
ssl_model_config = DenoisingAutoEncoderConfig(
    noise_strategy="swap",  # Can be 'zero' as well. Defines if noise is swapping features or making features zero
    default_noise_probability=0.7,  # Probability of corruption by noise
    include_input_features_inference=True,  # Setting this to True will include the input features in the concatenated layers in the decoder
    encoder_config=encoder_config,
    decoder_config=decoder_config,
    learning_rate=1e-3,
)

ssl_tabular_model = TabularModel(
    data_config=ssl_data_config,
    model_config=ssl_model_config,
    optimizer_config=ssl_optimizer_config,
    trainer_config=ssl_trainer_config,
    verbose=False,
)
# Pretraining
ssl_tabular_model.pretrain(train=ssl_train, validation=ssl_val)

print(ssl_tabular_model)


# Save model function
def save_model(model, path):
    model.save_model(path)


save_model(ssl_tabular_model, "SSL_model")

ssl_tabular_model_loaded = TabularModel.load_model("SSL_model")

print(ssl_tabular_model_loaded)

Maybe there is something different with the data you use..? if you can give me a Minimal Reproducible Example, I can take a look at it

[Draghont]

Hi,

It doesn't seem to be an issue with the data. From your code, it looks like you're saving the model after pretraining, but when you load it, you're not actually using the loaded model for create_finetune_model.

The error occurs after loading the model and when starting the fine-tuning process.

I managed to resolve a similar issue by checkpointing the pretrained model, loading it again, and then fine-tuning it. I used PyTorch's torch.save() and torch.load() methods for this:

Copy code
# Save the pretrained model
torch.save(model, path)

# Load the pretrained model
model = torch.load(path)

Using these methods, I was able to save the partially pretrained model and then proceed with fine-tuning successfully.

[manujosephv]

I think I have figured out the issue and have merged a PR fixing this. Maybe you can checkout latest GitHub source and check if it okay.

[manujosephv]

I hope this is not an issue anymore and the fix worked. Feel free to reopen if not.