TypeError: forward() got an unexpected keyword argument 'baseline value'. How do I correctly load a saved model in Skorch?

Question

I saved my Skorch neural net model using the below code:

net_b = NeuralNetClassifier(
    Classifier_b,
    max_epochs=50,
    optimizer__momentum= 0.9,
    lr=0.1,
    device=device,
)

#Fit the model on the full data
net_b.fit(merged_X_train, merged_Y_train);

#Test saving
import pickle
with open('MLP.pkl', 'wb') as f:
    pickle.dump(net_b, f)

When I try to load this model again and run it against test data, I receive the following error:

TypeError: forward() got an unexpected keyword argument 'baseline value'

This is my code:

#Split the data
X_train, y_train, X_valid, y_valid,X_test, y_test = train_valid_test_split(rescaled_data, target = 'fetal_health',
                                        train_size=0.8, valid_size=0.1, test_size=0.1)

input_dim = f_df_toscale.shape[1]
output_dim = len(np.unique(f_target))
hidden_dim_a = 20
hidden_dim_b = 12
device = 'cpu'

class Classifier_b(nn.Module):
    def __init__(self,
                 input_dim = input_dim,
                 hidden_dim_a = hidden_dim_b,
                 output_dim = output_dim):
        
        super(Classifier_b, self).__init__()

        #Take the inputs and pass these to a hidden layer
        self.hidden = nn.Linear(input_dim,hidden_dim_b)
        
        #Take the hidden layer and pass it through an additional hidden layer
        self.hidden_b = nn.Linear(hidden_dim_a,hidden_dim_b)
        
        #Take the hidden layer and pass to a multi nerouon output
        self.output = nn.Linear(hidden_dim_b,output_dim)

    def forward(self, x):
        hidden = F.relu(self.hidden(x))
        hidden = F.relu(self.hidden_b(hidden))
        output = F.softmax(self.output(hidden))     
        return output

#load the model
with open('MLP.pkl', 'rb') as f:
    model_MLP = pickle.load(f)

#Test the model
y_pred = model_MLP.predict(X_test)
ML = accuracy_score(y_test, y_pred)
print('The accuracy score for the MLP is ', ML)

When I run this model normally in the original notebook everything run fines. But when I try to load my model from a saved state I get the error. Any idea why? I have nothing called 'baseline value'.

Thanks

Answer 1

The save and load model can be problematic if the code changes. So it is better to use

save_params() and load_params()

In your case

net_b.save_params(f_params='some-file.pkl')

To load the model first initialize (initializing is very important) and then load parameters

new_net.initialize()

new_net.load_params(f_params='some-file.pkl')