526互联

关于保存自己的权重参数

发布时间 2023-07-24 17:37:10作者: Ray-Z

关于保存自己的权重参数

有的模型自己可以保存权重模型文件，那如果没有自己该怎么保存呢？

首先我们可以先查看一下，人家自带的权重模型文件，一般是 .pt 或 .pth 的文件，运行以下代码：

import torch

my_weights = torch.load(r'权重文件地址')
   
print('len=', len(my_weights.keys()))          
print('keys():', my_weights.keys())

image-20230724171402036

这是我的权重文件的，有3个键值对，关于 查看 ‘state_dict’ 的部分内容的代码：

print('len= ', len(my_weights['state_dict'].keys()))
print("\n".join(my_weights['state_dict'].keys()))

可以看到， ‘state_dict’ 中有159个键值对。

image-20230724171756738

而 epoch 和 best_acc 的结果是：

image-20230724171917604

这就是你训练模型的权重文件里面的内容了，每个模型的权重文件都不一样，需要自己根据自己模型自带的权重文件进行更改！！！

下面是我保存权重模型的代码：

my_weight = {
        'epoch': args.max_epochs,  #找到你设置epoch的变量名
        'best_acc': val_acc_max,  #同样，也是你模型中的值
        'state_dict': model.state_dict() #主要就是找到这个，训练的参数都在里面！！！
    }

    torch.save(
        my_weight, os.path.join(保存的地址, "my_weight.pt")  #保存模型！！！！
    )
    print("Model Was Saved ! ")  #打印输出提醒

这样，我们自己的权重模型文件就保存好啦！！！后面就可以用自己的权重模型文件啦！！！