pytorch网络训练结果保存后加载使用

先保存训练完后的参数

torch.save(mlp2.state_dict(), "./data/chap3/mlp2_param.pkl")

预处理数据

ss = StandardScaler(with_mean=True, with_std=True)
boston_Xs = ss.fit_transform(boston_X)
train_xt = torch.from_numpy(boston_Xs.astype(np.float32))
train_yt = torch.from_numpy(boston_y.astype(np.float32))
train_data = Data.TensorDataset(train_xt, train_yt)
train_loader = Data.DataLoader(
    dataset=train_data,
    batch_size=128,
    shuffle=True,
    num_workers=0
)

搭建一个跟之前相同的网络

class MLPmodel2(nn.Module):
    def __init__(self):
        super(MLPmodel2, self).__init__()
        self.hidden = nn.Sequential(
            nn.Linear(13, 10),
            nn.ReLU(),
            nn.Linear(10, 10),
            nn.ReLU()
        )
        self.regression = nn.Linear(10, 1)

    def forward(self, x):
        x = self.hidden(x)
        output = self.regression(x)
        return output

加载产生结果

modelx = MLPmodel2()#
modelx.load_state_dict(torch.load("./data/chap3/mlp2_param.pkl"))
output = modelx(train_xt).flatten()
loss_func = nn.MSELoss()
loss = loss_func(output, train_yt)
print(loss.item())