线性回归
import torch
import matplotlib.pyplot as plt
#生成一组数据
torch.manual_seed(100)
def get_fake_data(batch_size=8):
x=torch.rand(batch_size,1)*10
y=x*2+3
return x,y
#网络模型
class LinearModel(torch.nn.Module):
def __init__(self):
super(LinearModel, self).__init__()
self.linear = torch.nn.Linear(1, 1)
def forward(self, x):
y_pred = self.linear(x)
return y_pred
model = LinearModel()
#损失函数
criterion = torch.nn.MSELoss(size_average=False)
#优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
#训练神经网络
def train(x_data):
for epoch in range(1000):
y_pred = model(x_data)
#计算损失
loss = criterion(y_pred, y_data)
#梯度清0
optimizer.zero_grad()
#反向传播
loss.backward()
#参数更新
optimizer.step()
if __name__ == '__main__':
x_data, y_data = get_fake_data(batch_size=20)
train(x_data)
print('w = ', model.linear.weight.item())
print('b = ', model.linear.bias.item())
plt.scatter(x_data.numpy(), y_data.numpy())
plt.show()
plt.pause(0.5)