解决思路
由于我想是使用自己训练好的模型增加轮数继续训练,开始我将代码指向模型的地方在model_zoo.py中
if model_dir is None:
"""
如果model_dir没有指定(即None),则尝试寻找模型目录。
首先检查环境变量TORCH_HOME,如果没有设置,则默认使用~/.torch。
接着检查环境变量TORCH_MODEL_ZOO,如果没有设置,则默认使用TORCH_HOME(或默认的~/.torch)下的models目录作为模型目录。
"""
torch_home = os.path.expanduser(os.getenv("TORCH_HOME", "~/.torch"))
model_dir = os.getenv("TORCH_MODEL_ZOO", os.path.join(torch_home, "models"))
if not os.path.exists(model_dir):
os.makedirs(model_dir)
parts = urlparse(url)
filename = os.path.basename(parts.path)
if filename == "model_final.pkl":
# workaround as pre-trained Caffe2 models from Detectron have all the same filename
# so make the full path the filename by replacing / with _
filename = parts.path.replace("/", "_")
cached_file = os.path.join(model_dir, filename)
可以自己修改,我没有修改,我将模型训练好的名称model_final.pth改为vgg16_reducedfc.pth放到了C:\Users\Administrator.torch\models下,并执行train.py文件,但是报错了:RuntimeError: Error(s) in loading state_dict for ModuleList: Missing key(s)
解决问题
在vgg.py中修改load_state_dict函数即可
def init_from_pretrain(self, state_dict):
# 调用的load_state_dict方法用预训练的模型参数来初始化自己定义的新网络结构,pytorch中通用的用一个模型的参数初始化另一个模型的层的操作
# 重要的一个参数strict,默认为True,表示预训练的层和自己定义的网络结构严格对应相等,比如层名和维度
# self.vgg.load_state_dict(state_dict)
self.vgg.load_state_dict(state_dict,strict=False)
# 如果是加载本地下载好的预训练模型,就用torch.load(),再将strict改为False
# self.vgg.load_state_dict(torch.load(r"C:\Users\Administrator\.torch\models\vgg16_reducedfc.pth"),strict=False)
参考
https://huaweicloud.csdn.net/6380837bdacf622b8df8955d.html