SunnyHaze / IML-ViT

Official repository of paper “IML-ViT: Benchmarking Image manipulation localization by Vision Transformer”
MIT License
184 stars 23 forks source link

运行main_train.py报错(--resume读取release的测试用checkpoint导致的问题) #5

Closed CongYep closed 8 months ago

CongYep commented 8 months ago

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'

CongYep commented 8 months ago

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

SunnyHaze commented 8 months ago

感谢你的关注!

这个错误应该来源于你加了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中来测试你想要的数据集: https://github.com/SunnyHaze/IML-ViT/blob/912ebb291cd35dd960ac084c6d617f18639ac9cc/engine_train.py#L92

希望能够帮到你。

CongYep commented 8 months ago

好的,感谢