.Module.state_dict()2.1基本使用2.2保存和加载状态字典

1.保存、加载模型

torch.save()用于保存一个序列化对象到磁盘上该序列化对象可以是任何类型的对象包括模型、张量和字典等内部使用pickle模块实现对象的序列化。

数据会被保存为.pt、.pth格式可通过torch.load()从磁盘加载被保存的序列化对象加载时会重新构造出原来的对象。

1.保存整个模型继承了torch.nn.Module的类不推荐使用。

torch.load()利用pickle将保存的序列化对象反序列化得到原始数据。

可用于加载完整模型或状态字典。

#保存整个模型

torch.load(PATH)2.仅保存模型的参数状态字典state_dict推荐使用。

torch.nn.Module.load_state_dict()通过反序列化得到模型的state_dict()状态字典来加载模型传入的参数是状态字典而非.pt、.pth文件。

#只保存模型参数

model.load_state_dict(torch.load(PATH))在实际使用中推荐第二种方式第一种方式往往容易产生各种错误

设备错误。

若在cuda:0上训练好一个模型并保存则读取出来的模型也是默认在cuda:0上如果训练过程的其他数据被放到了cuda:1上那么就会发生错误

RuntimeError:

/opt/conda/conda-bld/pytorch_1503966894950/work/torch/lib/THC/generated/../generic/THCTensorMathPointwise.cu:215此时需要将其他其他数据都保存在cuda:0上或加载模型时指定使用cuda:1

device

map_locationdevice)版本错误比如使用pytorch1.0训练并保存CNN模型再用pytorch1.1读取模型则会出现错误

AttributeError:

padding_mode此时只能通过获取该模型的参数来加载新的模型

#加载模型参数

torch.load(model_path).state_dict()

#初始化新模型并加载参数

model.load_state_dict(model_state)2.torch.nn.Module.state_dict()

2.1基本使用

torch.nn.Module.state_dict()用于返回模型的状态字典其中保存了模型的可学习参数。

其中只有可学习参数的层卷积层、全连接层等和注册缓冲区batchnorm’s

running_mean才会作为模型参数保存优化器也有状态字典也可进行保存。

【例子】

__init__(self):super(TheModelClass,

nn.Conv2d(3,

self.pool(F.relu(self.conv1(x)))x

self.pool(F.relu(self.conv2(x)))x

x.view(-1,

model.state_dict():print(param_tensor,

\t,

model.state_dict()[param_tensor].size())#

print(Optimizers

optimizer.state_dict():print(var_name,

\t,

optimizer.state_dict()[var_name])查看模型与优化器的状态字典

2.2保存和加载状态字典

通过torch.save()来保存模型的状态字典state_dict即只保存学习到的模型参数并通过torch.nn.Module.load_state_dict()来加载并恢复模型参数。

PyTorch中最常见的模型保存扩展名为.pt或.pth。

#保存模型状态字典

model.load_state_dict(torch.load(PATH))

model.eval()

model.state_dict():print(param_tensor,

\t,

model.state_dict()[param_tensor].size())注意模型推理之前需要调用model.eval()函数将dropout和batch

normalization层设置为评估模式否则会导致模型推理结果不一致。

3.创建Checkpoint

模型检查点checkpoint是指模型训练过程中保存的模型状态包括模型参数权重与偏置、优化器状态等其他相关的训练信息。

通过保存检查点可以实现在训练过程中定期保存模型的当前状态以便在需要时恢复训练或用于模型评估和推理。

模型检查点常见的保存信息如下

1.模型权重模型的状态字典。

2.优化器状态优化器的状态字典。

3.训练状态当前的训练轮数epoch、批次batch等。

4.其他数据如学习率调度器的状态、自定义指标等。

【保存检查点】

model.state_dict(),optimizer_state_dict:

}#保存检查点

model.load_state_dict(checkpoint[model_state_dict])

optimizer.load_state_dict(checkpoint[optimizer_state_dict])#

恢复训练状态

__init__(self):super(SimpleModel,

nn.Linear(10,

y)loss.backward()optimizer.step()#

每10个epoch保存一次检查点if

model.state_dict(),optimizer_state_dict:

epoch,loss:

loss.item()}torch.save(checkpoint,

fcheckpoint_epoch_{epoch}.pth)#

加载检查点并继续训练

torch.load(checkpoint_epoch_10.pth)

model.load_state_dict(checkpoint[model_state_dict])

optimizer.load_state_dict(checkpoint[optimizer_state_dict])

start_epoch