GAT Missing/Unexpected key(s) in state_dict

[evantianzg]

I am using GAT for graph classification. I stored the model parameters of last several epochs through state_dict function. My GAT is defined as follows:

class GAT(torch.nn.Module):
    def __init__(self, num_features, dim_h, num_classes, heads=8):
        super(GAT, self).__init__()
        torch.manual_seed(2024)
        self.gat1 = GATConv(num_features, dim_h, heads=heads)
        self.gat2 = GATConv(dim_h*heads, dim_h, heads=1)
        self.lin = Linear(dim_h, num_classes)  # Add a linear layer for graph classification

    def forward(self, x, edge_index, batch):
        h = F.dropout(x, p=0.5, training=self.training)
        h = self.gat1(x, edge_index)
        h = F.elu(h)
        h = F.dropout(h, p=0.5, training=self.training)
        h = self.gat2(h, edge_index)
        h = global_mean_pool(h, batch)  # Pooling layer for graph classification
        h = self.lin(h)
        return h

I use the models of last several epochs to evaluate the test datasets and then average the performance. However, I got some trackback as follows:

Traceback (most recent call last):
  File "/Users/evan/workspace/codes/4gpu/conv.py", line 478, in <module>
    avg_performance_avg['precision'].append(avg_precision_avg)

  File "/Users/evan/workspace/codes/4gpu/conv.py", line 399, in main
    logger.info('\n')

  File "/Users/evan/workspace/codes/4gpu/conv.py", line 297, in evaluate_average
    y_pred = np.concatenate(y_pred)

  File "/Users/evan/workspace/codes/4gpu/conv.py", line 260, in evaluate
    # Return the last conv_epochs model weights file paths for averaging
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/evan/anaconda3/envs/pyt2.2/lib/python3.12/site-packages/torch/nn/modules/module.py", line 2153, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for GAT:
	Missing key(s) in state_dict: "gat1.att_src", "gat1.att_dst", "gat1.lin.weight", "gat2.att_src", "gat2.att_dst", "gat2.lin.weight".
	Unexpected key(s) in state_dict: "gat1.att", "gat1.lin_l.weight", "gat1.lin_l.bias", "gat1.lin_r.weight", "gat1.lin_r.bias", "gat2.att", "gat2.lin_l.weight", "gat2.lin_l.bias", "gat2.lin_r.weight", "gat2.lin_r.bias". 

I am not sure if it is the version problem or something wrong in my GAT model or the averaging performance implementation. What puzzles me is that this problem did not appear in the first place. It successfully executed several datasets, but the problem appeared after successfully executing several datasets.

Any suggestions are appreciated.🙏

[Romain-Nicolle]

Hi,

From what I see from your code, it seems to me that the problem lies somewhere where you get/load the parameters from the best_model.
I guess that you are using some kind of early stoppping criterion ? Like when the validation_loss stop decreasing after X epochs you stop and load the best_state ? Something like that ?
And you are saving the intermediate states ? Be careful to use a proper copy() of the model_state_dict.
If yes, how did you write this part of the code ? I guess it should look like something like that:

import copy
...
model_state_dict = copy.deepcopy(model.state_dict())
...
model.load_state_dict(model_state_dict)
...

You can only load_weights to a model that has been initialized if you are using lazy_initialization. I had a similar problem not so long ago and I had to adjust my code in order for it to work. I hope that your problem is similar and that it will solve it.

[akihironitta]

@evantianzg Can you share how model_weights is created?

[evantianzg]

I just initialize the model_weights in train function like this:

define train()

    last_model_weights = []  # Initialize the last model weights for each of the last conv_epochs

I want to save the models of last several consecutive epochs, so I use a list last_model_weights to achieve this.

[Romain-Nicolle]

Ok, so, just to be sure, your last_model_weights is a list containing several different dictionnaries with different model_state_dict. Okay got it. But according to this code snippet, there is no variable simply called model_weights.

In your evaluate() function, can you try to pass an argument like an index to get a precise dictionnary ? Something like:

def evaluate(index_of_the_model_to_get):
    model.load_state_dict(last_model_weights[index_of_the_model_to_get])
...

Ot maybe something like:

for key in best_models.keys():
    model.load_state_dict(best_models[key])
    # do stuff
    ...

[akihironitta]

model_weights should be a dict instead of list here, right?

def evaluate()

    model.load_state_dict(model_weights)

https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.load_state_dict

I see you export a state dict at:

best_model_wts_last_epoch = copy.deepcopy(model.state_dict())

but it's still unclear to me what the model_weights look like.

[evantianzg]

No, there is no variable called model_weights. model_weights only exists in my evaluate() function. I use best_models corresponding to different criteria to act as the specific best model in my train() function. It returns best_models after training.

def train()
    ......
    # Return the best models for each strategy
    best_models = {
        'last-epo': best_model_wts_last_epoch,
        'min-loss': best_model_wts_test_loss,
        'max-acc': best_model_wts_test_accuracy
    }

That allows me not to pass the index in the evaluate() function. I do not pass best_models in my evaluate() function neither. My evaluate() function looks like this:

def evaluate(model, loader, device, model_weights):
    model.load_state_dict(model_weights)
    model.eval()
    ......
    return accuracy, precision, recall, f1

best_models is used in my main() function, like this:

def main():
    best_model_wts = None  # Initialize the best model weights
            # Load the best model weights from the previous iteration if they exist
            if best_model_wts is not None:
                model.load_state_dict(best_model_wts)
                model, best_models = train(model, train_loader, test_loaders[-1], optimizer, criterion, ......)
                # Evaluate for each strategy for Single Best Model
                for strategy in strategies:
                    for i, test_loader in enumerate(test_loaders):
                        accuracy, precision, recall, f1 = evaluate(model, test_loader, device, best_models[strategy])
                        ......

The codes are indeed a little confusing. What puzzles me is that I split the dataset into multiple sub-datasets, and I use several types of models to handle these sub-datasets. It works well for some models (GCN, GraphSAGE, GIN) on all the sub-datasets, but it throw errors when using GAT on a specific sub-dataset. Although I am not sure if there is something wrong with my codes. I have removed all the early stopping part and strategies to choose best models, it works pretty well.