示例：模型定义与训练

以下示例展示了如何逐步搭建两层神经网络模型，并通过几行Python代码从头开始进行模型训练。

  1. import jittor as jt
  2. import numpy as np
  3. from jittor import nn, Module, init

以下代码定义了我们的模型，该模型是一个两层神经网络。 隐藏层的大小为10，激活函数为relu。

  1. ### model define
  3. class Model(Module):
  4.     def __init__(self):
  5.         self.layer1 = nn.Linear(1, 10)
  6.         self.relu = nn.ReLU()
  7.         self.layer2 = nn.Linear(10, 1)
  8.     def execute (self,x) :
  9.         x = self.layer1(x)
  10.         x = self.relu(x)
  11.         x = self.layer2(x)
  12.         return x

最后，从头开始训练该模型。 优化器使用简单的梯度下降，损失函数为L2距离。 为提高效率训练过程是异步的。 jittor通过统一计算图和即时分析器计算梯度，并进行计算图级和算子级的优化。

在该示例中，Jittor使用了多个优化，包括：算子融合，激活函数和损失函数可以融合到第一和第二全连接层中； 矩阵乘法中的三元算子也可以融合。 并行化，它可以提高现代多核CPU和GPU上计算密集型运算的性能。 算子融合是一种计算图级优化，而并行化则同时作用于图形级和算子级的优化。

  1. np.random.seed(0)
  2. jt.set_seed(3)
  3. n = 1000
  4. batch_size = 50
  6. def get_data(n):
  7.     for i in range(n):
  8.         x = np.random.rand(batch_size, 1)
  9.         y = x*x
  10.         yield jt.float32(x), jt.float32(y)
  12. model = Model()
  13. learning_rate = 0.1
  14. optim = nn.SGD (model.parameters(), learning_rate)
  16. for i,(x,y) in enumerate(get_data(n)):
  17.     pred_y = model(x)
  18.     loss = jt.sqr(pred_y - y)
  19.     loss_mean = loss.mean()
  20.     optim.step (loss_mean)
  21.     print(f"step {i}, loss = {loss_mean.data.sum()}")
  23. assert loss_mean.data < 0.005