Model behaves differently after saving and loading
I want to use torch.save() to save a trained model for inference. However, with either torch.load_state_dict() or torch.load(), I can't get the saved model. The loss computed by the loaded model is just different from the loss computed by the saved model. 

The relevant Libraries:

import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.nn import functional as F


The model:

class nn_block(nn.Module):

    def __init__(self, feats_dim):

        super(nn_block, self).__init__()

        self.linear = nn.Linear(feats_dim, feats_dim)
        self.bn = nn.BatchNorm1d(feats_dim)
        self.softplus1 = nn.Softplus()
        self.softplus2 = nn.Softplus()

    def forward(self, rep_mat):

        transformed_mat = self.linear(rep_mat)
        transformed_mat = self.bn(transformed_mat)
        transformed_mat = self.softplus1(transformed_mat)
        transformed_mat = self.softplus2(transformed_mat + rep_mat)

        return transformed_mat


class test_nn(nn.Module):

    def __init__(self, in_feats, feats_dim, num_conv, num_classes):

        super(test_nn, self).__init__()

        self.linear1 = nn.Linear(in_feats, feats_dim)
        self.convs = [nn_block(feats_dim) for _ in range(num_conv)]
        self.linear2 = nn.Linear(feats_dim, num_classes)
        self.softmax = nn.Softmax()

    def forward(self, rep_mat):

        h = self.linear1(rep_mat)
        for conv_func in self.convs:
            h = conv_func(h)
        h = self.linear2(h)
        h = self.softmax(h)

        return h


Train, save, and reload a model:

# fake a classification task
num_classes = 2; input_dim = 8
one = np.random.multivariate_normal(np.zeros(input_dim),np.eye(input_dim),20)
two = np.random.multivariate_normal(np.ones(input_dim),np.eye(input_dim),20)
inputs = np.concatenate([one, two], axis=0)
labels = np.concatenate([np.zeros(20), np.ones(20)])

inputs = Variable(torch.Tensor(inputs))
labels = torch.LongTensor(labels)

# build a model
net = test_nn(input_dim, 5, 2, num_classes)
optimizer = torch.optim.Adam(net.parameters(), lr=0.01)
net.train()

losses = []
best_score = 1e10
for epoch in range(25):

    preds = net(inputs)
    loss = F.cross_entropy(preds, labels)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    state_dict = {'state_dict': net.state_dict()}

    if loss.item()-best_score&lt;-1e-4: 
        # save only parameters 
        torch.save(state_dict, 'model_params.torch')
        # save the whole model 
        torch.save(net, 'whole_model.torch')
    best_score = np.min([best_score, loss.item()])

    losses.append(loss.item())

net_params = test_nn(input_dim, 5, 2, num_classes)
net_params.load_state_dict(torch.load('model_params.torch')['state_dict'])
net_params.eval()
preds_params = net_params(inputs)
loss_params = F.cross_entropy(preds_params, labels)
print('reloaded params %.4f %.4f' % (loss_params.item(), np.min(losses)))

net_whole = torch.load('whole_model.torch')
net_whole.eval()
preds_whole = net_whole(inputs)
loss_whole = F.cross_entropy(preds_whole, labels)
print('reloaded whole %.4f %.4f' % (loss_whole.item(), np.min(losses)))


As you can see by running the code, the losses computed by the two loaded models are different, while the two loaded models are exactly the same. Not just the two losses are different, they are also different from the loss computed by the best model that was saved in the first place.

Why this can happen? 

The state dict contains every parameter (nn.Parameter
) and buffer (similar to parameter, but which should not be trained/optimised) that has been registered on the module and all of its submodules. Everything else will not be included in that state dict.

Your test_nn module uses a list for convs, therefore it is not included in the state dict:

self.convs = [nn_block(feats_dim) for _ in range(num_conv)]


Not only are they not contained in the state dict, they are also not visible to net.parameters(), which means they are not trained/optimised at all.

To register the modules from the list you can wrap it in nn.ModuleList, which is a module that acts like a list, while correctly registering the modules it contains:

self.convs = nn.ModuleList([nn_block(feats_dim) for _ in range(num_conv)])


With that change both models produce the same result.

Since you are calling the convs modules sequentially in the for-loop (output of one module is the input of the next), you may consider using nn.Sequential, which you can call directly instead of having to use the for-loop. Sequencing is used a lot and it just makes it a little simpler, for example if you want to replace the sequence of modules with a single module, you don't need to change anything in the forward method.


  Not just the two losses are different, they are also different from the loss computed by the best model that was saved in the first place.


When you are training, you calculate the loss for the current input (batch) and then you optimise the parameters based on that input. This means your parameters differ from the ones used to calculate the loss. Because you are saving the model after that, it will also have a different loss (the one that would occur in the next iteration).

preds = net(inputs)
# Calculating the loss of the current model
loss = F.cross_entropy(preds, labels)

optimizer.zero_grad()
loss.backward()
# Updating the model's parameters based on the loss
optimizer.step()

# State of the model after it has been updated
state_dict = {'state_dict': net.state_dict()}

# Comparing the loss from BEFORE the update
# But saving the model from AFTER the update
if loss.item()-best_score&lt;-1e-4:
    # save only parameters
    torch.save(state_dict, 'model_params.torch')
    # save the whole model
    torch.save(net, 'whole_model.torch')


It's important to evaluate the model after the updates have been made. For this reason a validation set should be used, which is run after each epoch to assess the model's accuracy.
