PyTorch加载预训练模型

最新推荐文章于 2024-06-21 14:35:06 发布

xw2017

最新推荐文章于 2024-06-21 14:35:06 发布

阅读量884

点赞数 2

分类专栏： Python 文章标签： PyTorch 加载模型

本文链接：https://blog.csdn.net/xw2017/article/details/88862788

版权

Python 专栏收录该内容

13 篇文章 0 订阅

订阅专栏

加载单GPU模型

model = net()
pretrained_dict = torch.load("abc.pth")
model.load_state_dict(pretrained)

加载多GPU模型

model = net()
pretrained_dict = toch.load("m_abc.pth")
model.module.load_state_dict() # 多GPU要加module

加载部分预训练模型参数

model = net()
pretrained_dict = torch.load("abc.pth")
model_dict = model.state_dict()
# filter out unnecessary keys
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)