快速上手

本文将以训练 MNIST 数据集为例，简单地介绍如何使用 OneFlow 完成深度学习中的常见任务。通过文章中的链接可以跳转到各个子任务的专题介绍。

详细的介绍请阅读本文。让我们先从导入必要的库开始：

  1. import oneflow as flow
  2. import oneflow.nn as nn
  3. import oneflow.utils.vision.transforms as transforms
  5. BATCH_SIZE=128

加载数据

OneFlow 可以使用 Dataset 与 Dataloader 加载数据。

oneflow.utils.vision.datasets 模块中包含了不少真实的数据集(如 MNIST、CIFAR10、FashionMNIST)。

我们通过 oneflow.utils.vision.datasets.MNIST 获取 MNIST 的训练集和测试集数据。

  1. mnist_train = flow.utils.vision.datasets.MNIST(
  2.     root="data",
  3.     train=True,
  4.     transform=transforms.ToTensor(),
  5.     download=True,
  6.     source_url="https://oneflow-public.oss-cn-beijing.aliyuncs.com/datasets/mnist/MNIST/",
  7. )
  8. mnist_test = flow.utils.vision.datasets.MNIST(
  9.     root="data",
  10.     train=False,
  11.     transform=transforms.ToTensor(),
  12.     download=True,
  13.     source_url="https://oneflow-public.oss-cn-beijing.aliyuncs.com/datasets/mnist/MNIST/",
  14. )

输出：

  1. Downloading https://oneflow-public.oss-cn-beijing.aliyuncs.com/datasets/mnist/MNIST/train-images-idx3-ubyte.gz
  2. Downloading https://oneflow-public.oss-cn-beijing.aliyuncs.com/datasets/mnist/MNIST/train-images-idx3-ubyte.gz to data/MNIST/raw/train-images-idx3-ubyte.gz
  3. 9913344it [00:00, 36066177.85it/s]
  4. Extracting data/MNIST/raw/train-images-idx3-ubyte.gz to data/MNIST/raw
  5. ...

数据集下载并解压到 ./data 目录下。

利用 oneflow.utils.data.DataLoader 可以将 dataset 封装为迭代器，方便后续训练。

  1. train_iter = flow.utils.data.DataLoader(
  2.     mnist_train, BATCH_SIZE, shuffle=True
  3. )
  4. test_iter = flow.utils.data.DataLoader(
  5.     mnist_test, BATCH_SIZE, shuffle=False
  6. )
  8. for x, y in train_iter:
  9.     print("x.shape:", x.shape)
  10.     print("y.shape:", y.shape)
  11.     break

输出：

  1. x.shape: flow.Size([128, 1, 28, 28])
  2. y.shape: flow.Size([128])

🔗 Dataset 与 Dataloader

搭建网络

想要搭建网络，只需要实现一个继承自 nn.Module 的类就可以了。在它的 __init__ 方法中定义神经网络的结构，在它的 forward 方法中指定前向传播的计算逻辑。

  1. class NeuralNetwork(nn.Module):
  2.     def __init__(self):
  3.         super(NeuralNetwork, self).__init__()
  4.         self.flatten = nn.Flatten()
  5.         self.linear_relu_stack = nn.Sequential(
  6.             nn.Linear(28*28, 512),
  7.             nn.ReLU(),
  8.             nn.Linear(512, 512),
  9.             nn.ReLU(),
  10.             nn.Linear(512, 10),
  11.             nn.ReLU()
  12.         )
  14.     def forward(self, x):
  15.         x = self.flatten(x)
  16.         logits = self.linear_relu_stack(x)
  17.         return logits
  19. model = NeuralNetwork()
  20. print(model)

输出：

  1. NeuralNetwork(
  2.   (flatten): Flatten(start_dim=1, end_dim=-1)
  3.   (linear_relu_stack): Sequential(
  4.     (0): Linear(in_features=784, out_features=512, bias=True)
  5.     (1): ReLU()
  6.     (2): Linear(in_features=512, out_features=512, bias=True)
  7.     (3): ReLU()
  8.     (4): Linear(in_features=512, out_features=10, bias=True)
  9.     (5): ReLU()
  10.   )
  11. )

🔗 搭建神经网络

训练模型

为了训练模型，我们需要损失函数 loss_fn 和优化器 optimizer，损失函数用于评价神经网络预测的结果与 label 的差距；optimizer 调整网络的参数，使得网络预测的结果越来越接近 label（标准答案），这里选用 oneflow.optim.SGD。这一过程被称为反向传播。

  1. loss_fn = nn.CrossEntropyLoss()
  2. optimizer = flow.optim.SGD(model.parameters(), lr=1e-3)

定义一个 train 函数进行训练，完成前向传播、计算 loss、反向传播更新模型参数等工作。

  1. def train(iter, model, loss_fn, optimizer):
  2.     size = len(iter.dataset)
  3.     for batch, (x, y) in enumerate(iter):
  4.         # Compute prediction error
  5.         pred = model(x)
  6.         loss = loss_fn(pred, y)
  8.         # Backpropagation
  9.         optimizer.zero_grad()
  10.         loss.backward()
  11.         optimizer.step()
  13.         current = batch * BATCH_SIZE
  14.         if batch % 100 == 0:
  15.             print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

同时，定义一个 test 函数，用于检验模型的精度：

  1. def test(iter, model, loss_fn):
  2.     size = len(iter.dataset)
  3.     num_batches = len(iter)
  4.     model.eval()
  5.     test_loss, correct = 0, 0
  6.     with flow.no_grad():
  7.         for x, y in iter:
  8.             pred = model(x)
  9.             test_loss += loss_fn(pred, y)
  10.             bool_value = (pred.argmax(1).to(dtype=flow.int64)==y)
  11.             correct += float(bool_value.sum().numpy())
  12.     test_loss /= num_batches
  13.     print("test_loss", test_loss, "num_batches ", num_batches)
  14.     correct /= size
  15.     print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}, Avg loss: {test_loss:>8f}")

然后可以开始训练，定义5轮 epoch，每训练完一个 epoch 都使用 test 来评估一下网络的精度：

  1. epochs = 5
  2. for t in range(epochs):
  3.     print(f"Epoch {t+1}\n-------------------------------")
  4.     train(train_iter, model, loss_fn, optimizer)
  5.     test(test_iter, model, loss_fn)
  6. print("Done!")

输出：

  1. loss: 2.299633
  2. loss: 2.303208
  3. loss: 2.298017
  4. loss: 2.297773
  5. loss: 2.294673
  6. loss: 2.295637
  7. Test Error:
  8.  Accuracy: 22.1%, Avg loss: 2.292105
  10. Epoch 2
  11. -------------------------------
  12. loss: 2.288640
  13. loss: 2.286367
  14. ...

🔗 自动求梯度 🔗 反向传播与 optimizer

保存与加载模型

调用 oneflow.save 可以保存模型。保存的模型可以通过 oneflow.load 加载，用于预测等工作。

  1. flow.save(model.state_dict(), "./model")

🔗 模型的加载与保存

交流 QQ 群

安装或使用过程遇到问题，欢迎入群与众多 OneFlow 爱好者共同讨论交流：

加 QQ 群 331883 或扫描二维码

OneFlow 技术交流

为正常使用来必力评论功能请激活JavaScript