一般保存为.pt
格式,保存模型使用:
torch.save(model, '保存位置')
加载模型使用:
model_load = torch.load('加载模型的位置')
完整代码
import torch
import torch.nn as nn
class LinearRegressionModel(nn.Module):
def __init__(self, input_shape, output_shape):
super(LinearRegressionModel, self).__init__()
self.linear = nn.Linear(input_shape, output_shape)
def forward(self, x):
out = self.linear(x)
return out
if __name__ == '__main__':
model = LinearRegressionModel(10, 1)
torch.save(model, 'my_linear_model.pt')
model_load = torch.load('my_linear_model.pt')