自定义数据集使用框架的线性回归方法对其进行拟合

import torch
import numpy as np
import torch.nn as nn

criterion = nn.MSELoss()

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).unsqueeze(1)
y_train = torch.tensor(y_data, dtype=torch.float32)

model = nn.Sequential(nn.Linear(1, 1))

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

epoch = 500
for n in range(1, epoch + 1):
    y_pred = model(x_train)
    loss = criterion(y_pred.squeeze(1), y_train)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if n % 10 == 0 or n == 1:
        print(f'Epoch: {n}, Loss: {loss.item()}')