SunnyHaze
commented
4 months ago 因为main_train.py得到的checkpoint是按照字典形式组织的，包括了model，optimizer等等很多参数，会导致整个checkpoint大小超过1GB。其形式大致如下：
{
"model": <state_dict of model>,
"optimizer": <state_dict of model>,
......
}为了节省大小，demo.ipynb默认只读取model参数，也就是说它期待的读取的文件类型就是model字段读取出来实现，为了验证，可以把model.load_state_dict()中的strict参数改为True，这样只有严格match才能继续运行；或者提前将checkpoint中的
这里是我使用的重新save checkpoint的样例`Python脚本：
import torch
model = torch.load("/mnt/data0/xiaochen/workspace/IML-VIT/output_dir/checkpoint-150.pth")
output = model['model']
torch.save(output, "checkpoint-150.pth")English version:
The issue is that the checkpoint obtained from main_train.py is organized in dictionary format, including many parameters such as model, optimizer, and so on. This leads to a size checkpoint over 1GB. Its format is roughly as follows:
{
"model": <state_dict of model>,
"optimizer": <state_dict of model>,
......
}To save space, demo.ipynb by default only reads the model parameters. In other words, it expects to read a PyTorch checkpoint file that only contains an object of torch.load() section in demo.ipynb to extract the model field from the dictionary in advance. For validation purposes, you can change the strict parameter in model.load_state_dict() to True, so that only strict matching can proceed, or you can save the
Here is an example Python script I used to re-save the checkpoint:
import torch
model = torch.load("/mnt/data0/xiaochen/workspace/IML-VIT/output_dir/checkpoint-150.pth")
output = model['model']
torch.save(output, "checkpoint-150.pth")
linohzz
作者您好，我采用CASIAV2进行训练，CASIAV1作为验证，使用第108轮的权重文件在demo.ipynb上可视化示例图。遇到了以下问题：
图片可视化效果几乎不可见，使用训练集中的数据进行demo,效果也如上。 这是我的训练log.txt
使用您提供的权重文件进行可视化demo时，效果都正常