search

youngwanLEE / vovnet-detectron2

[CVPR 2020] VoVNet backbone networks for detectron2
Other
374 stars 69 forks source link

Used pretrained model mask_rcnn_V_99_FPN_3x on custom dataset #9

Closed DINOUS11 closed 4 years ago

DINOUS11 commented 4 years ago

I got this error while training with custom dataset

RuntimeError: The size of tensor a (81) must match the size of tensor b (6) at non-singleton dimension 0 image

youngwanLEE commented 4 years ago

@DINOUS11 My checkpoint includes an optimizer that was used in the training step on COCO.

I recommend

  1. load the checkpoint.
  2. remove the optimizer key and values or save the only model in the new file.
DINOUS11 commented 4 years ago

@youngwanLEE how can i save the only model in new file, thanks so much :+1:

youngwanLEE commented 4 years ago

Briefly,

checkpoint_path = "path/to/the/mask_V_99_eSE_ms_3x.pth"
new_path = ''path/to/the/mask_V_99_eSE_ms_3x.pth"

checkpoint = torch.load(checkpoint_path)

new_model_state_dict = {}
for k, v in checkpoint.items():
    if "model" in k:
        new_model_state_dict[k] = v

torch.save(new_model_state_dict, new_path)
DINOUS11 commented 4 years ago

@youngwanLEE it runs perfectly now, very appreciate for your work :100: image