Cannot load the weights after training with custom data

[PhucPhamSy]

Hi, after train model with distributed training method. I cannot load the weight name "last_model_state.pth". I trained with model name "fasterrcnn_resnet50_fpn"

Traceback (most recent call last):
File "/raid/data/phuc/detection/fastercnn-pytorch-training-pipeline/test_load_weight.py", line 37, in
model.load_state_dict(checkpoint['model_state_dict'])
File "/home/ubuntu/anaconda3/envs/fastercnn/lib/python3.9/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 FasterRCNN:
Missing key(s) in state_dict:

[sovit-123]

May I know which version of PyTorch you are using?

[PhucPhamSy]

May I know which version of PyTorch you are using?

I am using pytorch 2.0.1

[sovit-123]

Can you please paste the command that you are using?

[PhucPhamSy]

1. For training:
   export CUDA_VISIBLE_DEVICES=0,1
   python -m torch.distributed.launch --nproc_per_node=2 --use_env train.py --data data_configs/smoke.yaml --epochs 100 --model fasterrcnn_resnet50_fpn --name smoke_training --batch 16
2. For inference:
   python inference.py --input data/inference_data/image_1.jpg --weights outputs/training/smoke_training/last_model_state.pth

[PhucPhamSy]

But when I trained with only 1 GPU (not distributed train). I can use the command for inference above. Could you check the inference command for distributed train?

[sovit-123]

Oh. Interesting. I need some time to debug in this case.

[PhucPhamSy]

Many tks. I will wait your reply!

[aosaltik]

Hi, thanks a lot for such a nice repo! Is there any update for this issue? It's not working for eval.py as well in that situation.

[sovit-123]

Hello. Apologies for the delayed response. Recently I have not been able to put much time into the project due to time related constraints. However, I plan to resume again soon. I will be putting up a project page to work on the relevant issues and also new features soon.

Happy to receive feature recommendations from today.

[prof-nuduls]

In the inference.py I changed the section

if args['weights'] is not None:
        checkpoint = torch.load(args['weights'], map_location=DEVICE)
        # If config file is not given, load from model dictionary.
        if data_configs is None:
            data_configs = True
            NUM_CLASSES = checkpoint['data']['NC']
            CLASSES = checkpoint['data']['CLASSES']
        try:
            print('Building from model name arguments...')
            build_model = create_model[str(args['model'])]
        except:
            build_model = create_model[checkpoint['model_name']]
        model = build_model(num_classes=NUM_CLASSES, coco_model=False)
        model.load_state_dict(checkpoint['model_state_dict'])
    model.to(DEVICE).eval()

to

if args['weights'] is not None:
    checkpoint = torch.load(args['weights'], map_location=DEVICE, weights_only=True)
    
    if 'model_state_dict' in checkpoint:
        state_dict = checkpoint['model_state_dict']
        if any(key.startswith('module.') for key in state_dict.keys()):
            state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
        if data_configs is None:
            data_configs = True
            NUM_CLASSES = checkpoint['data']['NC']
            CLASSES = checkpoint['data']['CLASSES']
        try:
            build_model = create_model[str(args['model'])]
        except KeyError:
            build_model = create_model[checkpoint['model_name']]
        model = build_model(num_classes=NUM_CLASSES, coco_model=False)
        model.load_state_dict(state_dict)
        model.to(DEVICE)
        model.eval()
    else:
        raise KeyError("Checkpoint does not contain 'model_state_dict'. Please check the weights file.")

and that worked for me