简单线性回归
y = 2*x + 1
import numpy as np
import torch
import torch.nn as nnclass 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 outx_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.shapey_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时，注意转行成tensorinputs = 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'))
