自定义数据集 使用pytorch框架实现逻辑回归并保存模型，然后保存模型后再加载模型进行预测

import torch
import numpy as np
import torch.nn as nn
from torch.utils.data import DataLoader,TensorDataset

data = np.array([[-0.5, 7.7],
                 [1.8, 98.5],
                 [0.9, 57.8],
                 [0.4, 39.2],
                 [-1.4, -15.7],
                 [-1.4, -37.3],
                 [-1.8, -49.1],
                 [1.5, 75.6],
                 [0.4, 34.0],
                 [0.8, 62.3]])

x_data = data[:, 0]
y_data = data[:, 1]

x_train = torch.tensor(x_data, dtype=torch.float32)
y_train = torch.tensor(y_data, dtype=torch.float32)

dataset = TensorDataset(x_train, y_train)

dataloader = DataLoader(dataset=dataset, batch_size=2, shuffle=True)

criterion = nn.MSELoss()

class LinearModule(nn.Module):
    def __init__(self):
        super(LinearModule, self).__init__()
        self.layers = nn.Linear(1, 1)

    def forward(self, x):
        x = self.layers(x)
        return x

model = LinearModule()

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

epoch = 500

for n in range(1,epoch+1):
    epoch_loss = 0
    for x_batch, y_batch in dataloader:
        y_prd = model(x_batch.unsqueeze(1))
        loss = criterion(y_prd.squeeze(1), y_batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        epoch_loss += loss

    avg_loss = epoch_loss / len(dataloader)
    if n % 10 == 0 or n == 1:
        print(f'epoch:{n},loss:{avg_loss}')
        torch.save(model, 'model.pth')

load_model = torch.load('model.pth')

load_model.eval()

x_test = torch.tensor([1.8],dtype=torch.float32)

with torch.no_grad():
    y_pre = load_model(x_test)
print(f'x_test{x_test},y_pre:{y_pre}')