0%

使用pytorch进行一次线性回归

代码运行于python 3.7 pytorch 1.2.0

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
import torch
import random
import matplotlib.pyplot as plt

x = X = torch.arange(0,20,0.005).reshape(-1,1)
Y = 3*X+2.5

# 对输入数据加入gauss噪声
for i in range(len(X)):
X[i] += random.gauss(0, 0.35)
Y[i] += random.gauss(0, 0.75)
plt.plot(X.tolist(), Y.tolist(), linestyle='', marker='.',zorder=10)

W = torch.nn.Linear(1,1)
#优化权重和偏置
if torch.cuda.is_available():
W.cuda(0)
X = X.cuda(0)
x = x.cuda(0)
Y = Y.cuda(0)
optimizer = torch.optim.Adam(W.parameters())
criterion = torch.nn.MSELoss()
log = open('log.txt', 'a')
for step in range(10000):
optimizer.zero_grad()
loss = criterion(W(X),Y)
loss.backward()
optimizer.step()
if step % 1000 == 0:
txt = 'step:{}, = {}, loss = {}'.format(step,list(W.parameters()),loss)
log.write(txt + '\n')
print(txt)
log.close()
y = W(x)
plt.plot(x.tolist(),y.tolist(),color = 'y', linewidth=1,zorder=30)
plt.savefig('save.png', dpi=500)

计算结果