Pytorch保存模型saveModel及使用tensorboardX

最新推荐文章于 2025-06-03 17:42:45 发布

原创最新推荐文章于 2025-06-03 17:42:45 发布 · 2k 阅读

1 ·

CC 4.0 BY-SA版权

本文为博主CSDN Rosefun96原创文章，转载请注明。

学习Pytorch 专栏收录该内容

18 篇文章

订阅专栏

本文详细介绍了三种使用PyTorch保存和加载模型的方法，包括保存整个模型、仅保存模型参数以及通过Zip文件保存模型参数和配置。此外，还提供了使用TensorBoardX进行模型训练监控的指南。

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

1. 模型保存

方法一：

save model:

model = Mymodel()#按照你的实际模型
checkpoint = {'model': model,
              'state_dict': model.state_dict(),
              'optimizer' : optimizer.state_dict()}

torch.save(checkpoint, ckpt_file)

load model:

def load_checkpoint(filepath):
    checkpoint = torch.load(filepath)
    model = checkpoint['model']
    model.load_state_dict(checkpoint['state_dict'])
    for parameter in model.parameters():
        parameter.requires_grad = False
    
    model.eval()
    
    return model
model = load_checkpoint('checkpoint.pth')
print(model)

方法二：

torch.save(model, PATH)
# Model class must be defined somewhere
model = torch.load(PATH)
model.eval()

方法三：

from pathlib import Path
import shutil
import zipfile
import io
def load_model(filepath):
	with zipfile.ZipFile(filepath) as z:
		with z.open("model_params.json") as f:
			loaded_params = json.load(f)
		with z.open("network.pt") as f:# 或者pkl格式
			buffer = io.BytesIO(f.read())
		saved_state_dict = torch.load(buffer)

		
def save_model(path):
	"""
	Saving model with two distinct files.
	: param path: e.g. "/model/models"
	"""
	# Create folder
	Path(path).mkdir(parents=True, exist_ok=True)

	shutil.make_archive(path, 'zip', path)
	print(path)
	shutil.rmtree(path)

save_model("./model/models")
load_model("./model/models.zip")