pytorch加载预训练权重存在的问题

CV_er

已于 2024-01-09 21:42:23 修改

阅读量534

点赞数 7

文章标签： pytorch 人工智能深度学习

于 2024-01-09 18:00:06 首次发布

本文链接：https://blog.youkuaiyun.com/matlab_xiaogen/article/details/135486498

版权

本文讲述了在使用PyTorch时如何加载预训练模型的权重，特别提到了处理模型结构变化时，如何通过`model.load_state_dict()`函数并调整`strict`参数，以及一个与模型加载无关但相关的Infer类实例化过程。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

weigths = torch.save(model,'save.pt')
model.load_state_dict(weigths, strict = True)  #model.load_state_dict()函数把加载的权重复制到模型的权重中去

这里面有个小坑，strict布尔值默认True，对，默认严格导入。如果对网络结构有所改变，则需要删除相应的权重删除，哎都是为了导入backbone的权重呀，说多了都是泪。

new_weigths = {}
for k, v in weigths.items():
    if "fc" in k:
        continue
    new_weigth[k] = v

model.load_state_dict(weigths, strict = False)

因为我改变了全连接层，所以需要删除相应的权重，再将 strict设置成Fasle

#pragma once
#include <string>

class Infer
{
private:
    /* data */
    Infer(const std::string& );
public:

    Infer(const Infer& instance) = delete;
    // （此static是修饰函数）静态成员函数返回静态值
    static Infer* get_instance(const std::string& onnx_path){
        // 局部静态变量声明是可以定义的，因为静态不属于该类，且多线