pytorch 保存模型 (.pt, .pth, .pkl)

.pt, .pth, .pkl 僅僅是后綴不同

一、只保存模型權重參數，不保存模型結構

保存：

torch.save(model.state_dict(), 'mymodel.pth')

加載：

torch.load('mymodel.pth路徑') #加載后是一個OrderedDict類型

調用：

model = My_model(*args, **kwargs)  #這里需要重建模型結構，My_model

model.load_state_dict(torch.load(mymodel.pth))

二、保存整個模型

保存：

torch.save(model, mymodel.pth)

加載：

model = torch.load('mymodel.pth路徑') #加載后是一個Class類型

可以通過  ⬇ 打印實例的所有屬性和方法

for item in dir(model):
　　print(item)

調用：

不需要建模型結構，就是model