91 lines
3.2 KiB
Python
91 lines
3.2 KiB
Python
from enum import IntEnum, auto
|
||
from pathlib import Path
|
||
import sys
|
||
import torch
|
||
import matplotlib.pyplot as plt
|
||
import torch.nn.functional as F
|
||
|
||
sys.path.append(str(Path(__file__).resolve().parent.parent))
|
||
import pytorch_gpu_utils
|
||
|
||
|
||
class CurveKind(IntEnum):
|
||
"""生成假数据时使用的曲线"""
|
||
Polynomials = auto()
|
||
Sine = auto()
|
||
|
||
|
||
class DataSource:
|
||
"""用于拟合的随机生成的假数据"""
|
||
|
||
x: torch.Tensor
|
||
y: torch.Tensor
|
||
|
||
def __init__(self, device: torch.device, curve_kind: CurveKind):
|
||
match curve_kind:
|
||
case CurveKind.Polynomials:
|
||
x = torch.linspace(-1, 1, steps=100).reshape(-1, 1)
|
||
y = -x.pow(3) + 2 * x.pow(2) + 0.2 * torch.rand(x.size())
|
||
case CurveKind.Sine:
|
||
# 正弦在0-2之间变化才不是类似线性的
|
||
x = torch.linspace(0, 2, steps=100).reshape(-1, 1)
|
||
y = x.sin() + 0.2 * torch.rand(x.size())
|
||
|
||
self.x = x.to(device)
|
||
self.y = y.to(device)
|
||
|
||
|
||
class Net(torch.nn.Module):
|
||
"""继承torch的module用于表示网络"""
|
||
|
||
def __init__(self, n_feature, n_hidden, n_output):
|
||
super(Net, self).__init__() #继承_init_功能
|
||
#定理每层用什么样的形式
|
||
self.hidden1 = torch.nn.Linear(n_feature, n_hidden) #隐藏层线性输出
|
||
self.hidden2 = torch.nn.Linear(n_hidden, n_hidden) #输出层线性输出
|
||
self.hidden3 = torch.nn.Linear(n_hidden, n_hidden) #输出层线性输出
|
||
self.predict = torch.nn.Linear(n_hidden, n_output) #输出层线性输出
|
||
|
||
def forward(self, x): #这同时也是module中的forward功能
|
||
#正向传播输入值，神经网络分析出输出值
|
||
x = F.relu(self.hidden1(x)) #激励函数（隐藏层的线性值）
|
||
x = F.relu(self.hidden2(x))
|
||
x = F.relu(self.hidden3(x))
|
||
x = self.predict(x) #输出值
|
||
return x
|
||
|
||
|
||
def main():
|
||
device = pytorch_gpu_utils.get_gpu_device()
|
||
test_data = DataSource(device, CurveKind.Polynomials)
|
||
net = Net(n_feature=1, n_hidden=20, n_output=1).to(device)
|
||
|
||
#optimizer是训练的工具
|
||
optimizer = torch.optim.SGD(net.parameters(), lr=0.01) #传入net的所有参数，学习率
|
||
loss_func = torch.nn.MSELoss() #预测值和真实值的误差计算公式（均方差）
|
||
|
||
for t in range(2000):
|
||
optimizer.zero_grad() #清空上一步的残余更新参数值
|
||
prediction: torch.Tensor = net(test_data.x) #喂给net训练数据x，输出预测值
|
||
loss: torch.Tensor = loss_func(prediction, test_data.y) #计算两者的误差
|
||
loss.backward() #误差反向传播，计算参数更新值
|
||
optimizer.step() #将参数更新值施加到net的parameters上
|
||
|
||
#plot and show learning process
|
||
plt.cla()
|
||
plt.scatter(test_data.x.cpu().data.numpy(), test_data.y.cpu().data.numpy())
|
||
plt.scatter(test_data.x.cpu().data.numpy(), prediction.cpu().data.numpy())
|
||
plt.text(0.5,
|
||
0,
|
||
'Loss=%.4f' % loss.cpu().data.numpy(),
|
||
fontdict={
|
||
'size': 20,
|
||
'color': 'red'
|
||
})
|
||
plt.show()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
pytorch_gpu_utils.print_gpu_availability()
|
||
main()
|