import jittor as jt #骨干网络模块 from jittor import Module #神经网络模块 from jittor import nn import numpy as np import matplotlib matplotlib.use('TkAgg') import matplotlib.pyplot as plt #线性回归实例 #该模型是一个两层神经网络。 隐藏层的大小为10，激活函数为relu class Model(Module): def __init__(self): self.layer1 = nn.Linear(1, 10) self.relu = nn.Relu() self.layer2 = nn.Linear(10, 1) def execute (self,x) : x = self.layer1(x) x = self.relu(x) x = self.layer2(x) return x def get_data(n): # generate random data for training test. for i in range(n):#n=1000 #产生一个numpy.ndarray数据类型的列表 #其中包含batch_size（50）个numpy.ndarray数据类型的列表x #每个小列表里只有一个（0，1）的数，数据类型是'numpy.float64,x[index] #y跟x一样 x = np.random.rand(batch_size, 1) y = x*x #返回一个generator实例 yield jt.float32(x), jt.float32(y) #随机数种子对后面的结果一直有影响,后面的随机数组都是按一定的顺序生成的 np.random.seed(0) jt.set_seed(3) n = 1000 batch_size = 50 #新建模型 model = Model() #设置学习效率 learning_rate = 0.1 #保持当前参数状态并基于计算得到的梯度进行参数更新 optim = nn.SGD(model.parameters(), learning_rate) min_loss = 1.0 #开启一个画图的窗口进入交互模式，实时更新数据 plt.ion() #优化器使用简单的梯度下降，损失函数为L2距离 #enumerate:将一个可遍历的数据对象(如列表、元组或字符串)组合为一个索引序列，同时列出数据和数据下标 #列表中的每个元素都是元组(x,y) for i,(x,y) in enumerate(get_data(n)): pred_y = model(x) loss = ((pred_y - y)**2) loss_mean = loss.mean() optim.step (loss_mean) #print(f"step {i}, loss = {loss_mean.data.sum()}") #根据每次的loss来选择是否绘图，保证最后一张图的loss最小 if(loss_mean.data[0]