Jittor线性回归+matplotlib图表动态展示

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]<min_loss):
        min_loss=loss_mean.data[0]
        #清除刷新前的图
        plt.clf()
        plt.suptitle(str(i)+"time loss:"+str(loss_mean.data),fontsize=10)
        #第一张图
        a_graph = plt.subplot(2,1,1)
        a_graph.set_title('Raw data(x,y)')
        a_graph.set_xlabel('x',fontsize=10)
        a_graph.set_ylabel('y',fontsize=10)
        plt.plot(x.data,y.data,'r^')
        #第二张图
        b_graph=plt.subplot(2,1,2)
        b_graph.set_title("Fitted data(x,pred_y)")
        b_graph.set_xlabel('x',fontsize=10)
        b_graph.set_ylabel('pred_y',fontsize=10)
        plt.plot(x.data,pred_y.data,'g-')
        #弹窗停留0.4秒
        plt.pause(0.4)

#关闭交互模式
plt.ioff()
plt.show()

结果：