定义网络

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

全部代码

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__':
    x_train = torch.randn(100, 4)  # 生成100个4维的随机数，作为训练集的 X
    y_train = torch.randn(100, 1)  # 作为训练集的label

    # 指定参数与损失函数
    model = LinearRegressionModel(x_train.shape[1], 1)  # 对应__init__方法的input_shape:输入的维度；output_shape:输出的维度
    epochs = 1000  # 迭代1000次
    learning_rate = 0.01  # 学习率
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)  # 优化函数
    criterion = nn.MSELoss()  # Loss使用MSE值，目标是使MSE最小

    for epoch in range(epochs):
        epoch += 1
        # 梯度要清零每一次迭代
        optimizer.zero_grad()
        # 前向传播
        outputs = model(x_train)
        # 计算损失
        loss = criterion(outputs, y_train)
        # 返向传播
        loss.backward()
        # 更新权重参数
        optimizer.step()
        # 每隔50次打印一次
        if epoch % 50 == 0:
            print('epoch {}, loss {}'.format(epoch, loss.item()))
    # 预测
    predicted = model(torch.randn(100, 4)).data.numpy()
    print(predicted)