torch神經網路--線性迴歸

星空28發表於2024-10-05
神經網路

簡單線性迴歸

y = 2*x + 1

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


class LinearRegressionModel(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(LinearRegressionModel, self).__init__()
        self.linear = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        out = self.linear(x)
        return out


x_values = [i for i in range(11)]
x_train = np.array(x_values, dtype=np.float32)
x_train = x_train.reshape(-1, 1)
x_train.shape

y_values = [2*i+1 for i in x_values]
y_train = np.array(y_values, dtype=np.float32)
y_train = y_train.reshape(-1, 1)
y_train.shape
input_dim = 1
output_dim = 1
model = LinearRegressionModel(input_dim, output_dim)

# 如果使用GPU訓練,增加以下兩行程式碼
# device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# model.to(device)


# 指定好引數和損失函式
epochs = 1000
learning_rate = 0.01
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
criterion = nn.MSELoss()

# 訓練模型
for epoch in range(epochs):
    epoch += 1
    # 使用cpu時，注意轉行成tensor
    inputs = torch.from_numpy(x_train)
    labels = torch.from_numpy(y_train)
    # 如果使用GPU訓練,將以上兩行程式碼修改為
    # inputs = torch.from_numpy(x_train).to(device)
    # labels = torch.from_numpy(y_train).to(device)

    # 梯度要清零每一次迭代
    optimizer.zero_grad()
    # 前向傳播
    outputs = model(inputs)
    # 計算損失
    loss = criterion(outputs, labels)
    # 反向傳播
    loss.backward()
    # 更新權重引數
    optimizer.step()

    # 列印
    if epoch % 50 == 0:
        print('epoch {}, loss {}'.format(epoch, loss.item()))


# CPU測試模型預測結果
predicted = model(torch.from_numpy(x_train).requires_grad_()).data.numpy()

# 模型的儲存
torch.save(model.state_dict(), 'model.pkl')
# 模型讀取
model.load_state_dict(torch.load('model.pkl'))

相關文章