pytorch model saving method (.pt, .pth, .pkl)

The difference between different suffixes of the model

I often see pytorch model files with suffixes of .pt, .pth, and .pkl . In fact, they are not different in format, but the suffixes are different (that’s all). Use the torch.save() function to save the model When it comes to files, everyone has different preferences. Some people like to use the .pt suffix, and some people like to use .pth or .pkl. The model files saved with the same torch.save() statement are no different.

In pytorch's official documentation/code, there are .pt and .pth. The general practice is to use .pth, but it seems that there are more .pt in the official documents, and the official does not care much about using one.

Model saving and calling method 1:

Only save model parameters, not model structure

keep:

torch.save(model.state_dict(), mymodel.pth)#只保存模型权重参数,不保存模型结构

transfer:

model = My_model(*args, **kwargs)  #这里需要重新模型结构,My_model
model.load_state_dict(torch.load(mymodel.pth))#这里根据模型结构,调用存储的模型参数
model.eval()

Model saving and calling method 2:

Save the entire model, including model structure + model parameters

keep:

torch.save(model, mymodel.pth)#保存整个model的状态

transfer:

model=torch.load(mymodel.pth)#这里已经不需要重构模型结构了,直接load就可以
model.eval()

Guess you like

Origin blog.csdn.net/Cretheego/article/details/128789192