pytorch保存训练好的模型

PyTorch提供了两种保存训练好的模型的方法。 
第一种是只保存模型参数,这也是推荐的方法:

#保存
torch.save(the_model.state_dict(), PATH)
#读取
the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))

第二种方法保存整个模型:

#保存
torch.save(the_model, PATH)
#读取
the_model = torch.load(PATH)

注:PATH的格式:'./model_file_name/the_model_name.tar'

发布了52 篇原创文章 · 获赞 6 · 访问量 9006

猜你喜欢

转载自blog.csdn.net/PMPWDF/article/details/101162137