原文地址:Pytorch模型中的parameter与buffer - 知乎 (zhihu.com)
保存原来模型的参数:
torch.save(model.state_dict(), PATH)
# load
model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()
model.state_dict()返回的是OrderedDict类型,包含了两种参数,分别是parameter和buffer
- parameter在训练的过程中是需要不断优化的
- 返回方式:model.parameters()
- 注册方式
- register_buffer(name, tensor) name: string tensor: Tensor
- register_parameter(name, param) name: string param: Parameter
- buffer是训练的不变量,在训练过程中不需要应用梯度更新
- 返回方式:model.buffers()
- 注册方式:register_buffer()