import torch
import torch.nn as nn
import numpy as np
x_values = [i for i in range(11)]
x_train = np.array(x_values, dtype=np.float32)
x_train = x_train.reshape(-1, 1)
y_values = [2 * i + 1 for i in range(11)]
y_train = np.array(y_values, dtype=np.float32)
y_train = y_train.reshape(-1, 1)
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
input_dim = 1
output_dim = 1
model = LinearRegressionModel(input_dim, output_dim)
print(type(model))
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
inputs = torch.from_numpy(x_train)
labels = torch.from_numpy(y_train)
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()) )
torch.save(model.state_dict(), 'model.pkl')
pytorch构建线性回归代码示例
于 2022-01-21 21:43:22 首次发布