SAM2UNet-MSD.pth chkpt file does not contain key model

[msudhanshu-nus]

When I am loading this kept in train.sh, running it is producing the following error:
Traceback (most recent call last):
File "/mnt/iMVR/sudhanshu/SAM2-UNet/train.py", line 87, in
main(args)
File "/mnt/iMVR/sudhanshu/SAM2-UNet/train.py", line 46, in main
model = SAM2UNet(args.hiera_path)
File "/mnt/iMVR/sudhanshu/SAM2-UNet/SAM2UNet.py", line 129, in init
model = build_sam2(model_cfg, checkpoint_path)
File "/mnt/iMVR/sudhanshu/SAM2-UNet/sam2/build_sam.py", line 36, in build_sam2
_load_checkpoint(model, ckpt_path)
File "/mnt/iMVR/sudhanshu/SAM2-UNet/sam2/build_sam.py", line 88, in _load_checkpoint
sd = torch.load(ckpt_path, map_location="cpu")["model"]
KeyError: 'model'

I checked the keys of this kept, but could not find key "model" in it, and therefore this error.
Kindly check if the kept file uploaded is alright.

Regards,
Sudhanshu.

[xiongxyowo]

Hi, if you want to train a new model starting from the original SAM2, the "args.hiera_path" should be filled with the checkpoint path of SAM2 itself, such as sam2_hiera_large.pt:

# train.py
args.hiera_path = "./sam2_hiera_large.pt"
model = SAM2UNet(args.hiera_path)

If you want to test the trained SAM2-UNet model, you can load the weights of SAM2-UNet, such as "SAM2UNet-MSD.pth":

# test.py
args.checkpoint = "./SAM2UNet-MSD.pth"
model = SAM2UNet()
model.load_state_dict(torch.load(args.checkpoint), strict=True)

[msudhanshu-nus]

I want to use SAM2UNet-MSD.pth to train on my dataset for fine-tuning purpose. Basically, the camaflouged features learned are already there in in SAM2UNet-MSD.pth and will be useful to me in further fine-tuning. Can it be done?

Basically, instead of sam2_hiera_large.pt, I want to put SAM2UNet-MSD.pth which is causing the issue.

train.py

args.hiera_path = "./SAM2UNet-MSD.pth"
model = SAM2UNet(args.hiera_path)

[xiongxyowo]

It is possible to use a pre-trained SAM2-UNet during the training phase. Note that when further fine-tuning on SAM2-UNet, the original weight of SAM2 is no longer necessary:

# train.py
# args.hiera_path = "./sam2_hiera_large.pt"
args.checkpoint = "./SAM2UNet-MSD.pth"
model = SAM2UNet()
model.load_state_dict(torch.load(args.checkpoint), strict=True)