pytorch神经网络模型的保存与读取

小六的画布

于 2022-07-14 21:02:12 发布

阅读量632

点赞数

CC 4.0 BY-SA版权

分类专栏： pytorch 文章标签：神经网络 pytorch 深度学习

本文链接：https://blog.youkuaiyun.com/qq_42202395/article/details/125792603

pytorch 专栏收录该内容

28 篇文章

订阅专栏

本文详细介绍了如何使用PyTorch在两种方式下保存和读取神经网络模型，包括完整模型与参数分离的方法，以及可能遇到的陷阱和实例演示。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

pytorch神经网络模型的保存与读取

#model_save
vgg16 = torchvision.models.vgg16(pretrained = False)

#保存方式1
torch.save(vgg16, "vgg16_method1.pth") #保存模型和参数
#方式2， 官方推荐，更小
torch.save(vgg16.state_dict(),"vgg16_method2.pth")  #参数保存为字典，模型不再保存

#陷阱
class Model_test(nn,Module):
	def __init__():
		super(Model_test，self).__init__()
		self.conv1 = nn.Conv2d(3,64,kernel_size = 3)
	def forward(self, x):
		x = self.conv1(x)
		return x

#model_load
#方式1，加载是模型
vgg16 = torch.load("vgg16_method1.pth")
print(model)
#方式2，加载模型
vgg16 = torchvision.models.vgg16(pretrained = False)
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))
print(vgg16 )


#陷阱(需要copy过来)
class Model_test(nn,Module):
	def __init__():
		super(Model_test，self).__init__()
		self.conv1 = nn.Conv2d(3,64,kernel_size = 3)
	def forward(self, x):
		x = self.conv1(x)
		return x
model = torch.load(model_test_method1.pth)