模型参数常见读取与保存操作

简洁版torch.save和torch.load模型参数

	# save
    torch.save(model.state_dict(), "/home/yingmuzhi/microDL_2_0/a.pth")

	# load
	model = UNet()
    model.load_state_dict(torch.load("/home/yingmuzhi/microDL_2_0/a.pth"))

如何使用torch.save存储梯度信息

默认情况下，torch.save不会保存张量的梯度信息。它只会保存张量的值、形状和数据类型。

如果要保存梯度信息，可以在调用torch.save时将张量的requires_grad属性设置为True。这样，保存的文件中将包含梯度信息。

例如：

import torch

# 创建一个需要梯度的张量
x = torch.randn(3, 4, requires_grad=True)

# 计算一些操作
y = 2 * x + 1
z = y.mean()

# 计算梯度
z.backward()

# 保存张量和梯度信息
torch.save(x, 'tensor.pt')

# 加载张量和梯度信息
loaded_x = torch.load('tensor.pt')
print(loaded_x)

在这个例子中，x张量的梯度信息被保存在文件’tensor.pt’中。加载张量时，梯度信息也将被恢复。

模型参数导入存储模板

模型参数常见读取与保存操作，代码见下：

'''
经常存储的会有

- 模型参数, 包括权重weights和偏置bias, 即model.params.weight和model.params.bias
- optimizer
- epoch迭代次数
- ...

我们的处理流程往往是：
先用torch.load()导入, 读成checkpoint
将checkpoint分步骤解析, 如有的可能需要使用load_state_dict()来读取
将需要存储的东西以一个字典的形式组装成checkpoint, 其中可能有些需要存储的东西要使用state_dict()来存储
使用torch.save()将checkpoint存储

... 表示无实意的pass或者: 
参考`https://zhuanlan.zhihu.com/p/264896206`

'''
import os, torch


SAVE_DICTIONARY: dict = {
    "model": {},
    "optimizer": {},
    "start_epoch": 1,
    "args": [],
}



def main(

):
    ...
    
    resume = '',
    model = None,
    optimizer = None,
    # load pre-trained model
    if os.path.exists(resume):
        checkpoint = torch.load(resume, map_location="cpu") # 先用torch.load()全部导入, 读成checkpoint再后续分解
        model.load_state_dict(checkpoint["model"])          # 有些要存储成torch的字典形式的, 就必须使用.load_state_dict()来将dict解压 和 用.state_dict()来压缩成dict 
        optimizer.load_state_dict(checkpoint["optimizer"])
        # lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
        start_epoch = checkpoint["epoch"] + 1
        print("load pre-trained model successfully!")
    else:
        print("load pre-trained model failed.")
    
    ...

    epoch = 1
    args = []
    # save model
    checkpoint = {
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        # "lr_scheduler": lr_scheduler.state_dict(),
        "epoch": epoch, 
        "args": args
    }                                                       # 先组装成字典checkpoint, 再使用torch.save()全部存储
    torch.save(checkpoint, resume)

    ...