KeyError: 'net' when trying to run test

[borisandeva]

I was trying to run the test script with the command:

CUDA_VISIBLE_DEVICES=0 python -W ignore test.py --resume ./out/epoch_61.pth --data_root_path /mnt/zzhou82/PublicAbdominalData/ --store_result --cache_dataset --cache_rate 0.6

and got:

(universal) ~/CLIP-Driven-Universal-Model$ CUDA_VISIBLE_DEVICES=0 python -W ignore test.py --resume pretrained_weights/swin_unetr.base_5000ep_f48_lr2e-4_pretrained.pt --data_root_path /gpu_home/bori/CLIP-Driven-Universal-Model/data --store_result --cache_dataset --cache_rate 0.6 --dataset_list "pancreas_val"
Traceback (most recent call last):
  File "/gpu_home/bori/CLIP-Driven-Universal-Model/test.py", line 209, in <module>
    main()
  File "/gpu_home/bori/CLIP-Driven-Universal-Model/test.py", line 185, in main
    load_dict = checkpoint['net']
KeyError: 'net'

If i print out the keys:

ipdb> checkpoint.keys()
dict_keys(['epoch', 'best_acc', 'state_dict'])

I tried substituting the "net" with "state_dict" key but it fails.

Some info about my setup:

 OS: Ubuntu 22.04 jammy
 Kernel: x86_64 Linux 6.1.12-060112-generic
 Uptime: 39d 3h 50m
 Packages: 2087
 Shell: bash 5.1.16
 Disk: 2.7T / 10T (29%)
 CPU: 13th Gen Intel Core i7-13700KF @ 24x 5.3GHz [53.0°C]
 GPU: NVIDIA RTX A6000
 RAM: 5237MiB / 96377MiB

Python 3.10.13

[ljwztc]

you should resume weight in https://www.dropbox.com/s/jdsodw2vemsy8sz/swinunetr.pth, instead of pretrained_weights/swin_unetr.base_5000ep_f48_lr2e-4_pretrained.pt, which is used for training.

[ljwztc]

I have updated the inference instruction in README. You can find here.