pytorch权重保存与加载

pytorch权重保存:

# 1) 保存网络的结构以及权重值文件
torch.save(model,'XXX.pth')  
# 2) 仅保存权重值文件
torch.save(model.state_dict,'XXX.pth')

pytorch权重加载

# 1) 加载有结构有权重值的参数文件
model = torch.load('XXX.pth')
# 2) 加载只有权重值的的参数文件
model = model.load_state_dict(torch.load('XXX.pth'))

WHAT IS A STATE_DICT IN PYTORCH

In PyTorch, the learnable parameters (i.e. weights and biases) of a torch.nn.Module model are contained in the model’s parameters (accessed with model.parameters()). A state_dict is simply a Python dictionary object that maps each layer to its parameter tensor.

state_dict是一个python字典对象,包含了每一层网络的参数张量(parameter tensor)

根据pytorch,不仅有model.state_dict(),还有optimizer.state_dict()

你可能感兴趣的:(pytorch)