运行main_train.py报错（--resume读取release的测试用checkpoint导致的问题）

Traceback (most recent call last):
File "main_train.py", line 269, in
main(args)
File "main_train.py", line 208, in main
misc.load_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler)
File "/home/xxx/my_project/IML-ViT-main/utils/misc.py", line 324, in load_model
model_without_ddp.load_state_dict(checkpoint['model'])
KeyError: 'model'

还有一个问题，请问如何用你公布的iml-vit_checkpoint.pth进行测试

感谢你的关注！

这个错误应该来源于你加了resume参数用于“恢复训练”，然而可能由于resume是我们release的checkpoint导致报错。该checkpoint为了减少大小，没有包含optimizer的梯度等等信息。（如果包含的话会达到1个G）

如果真正要load的话，应该是要load一个词典，以如下形式组织：

{
   "model" : <model‘s state_dict>,
   "optimizer" : <optimizer‘s state_dict>,
   "loss_scaler" : <loss_scaler's state_dict>
}

目前你可以参考这个函数来测试，可以将其组装到目前提供的demo.py中来测试你想要的数据集：

IML-ViT/engine_train.py

Line 92 in 912ebb2

def test_one_epoch(model: torch.nn.Module,

希望能够帮到你。

好的，感谢