QUICKSTART

This section will take the training process of MNIST as an example to briefly show how OneFlow can be used to accomplish common tasks in deep learning. Refer to the links in each section to the presentation on each subtask.

Let’s start by importing the necessary libraries:

import oneflow as flow
import oneflow.nn as nn
import oneflow.utils.vision.transforms as transforms
BATCH_SIZE=128

Working with Data

OneFlow has two primitives to work with data, which are Dataset and Dataloader.

The oneflow.utils.vision.datasets module contains a number of real data sets (such as MNIST, CIFAR 10, FashionMNIST).

We can use oneflow.utils.vision.datasets.MNIST to get the training set and test set data of MNIST.

mnist_train = flow.utils.vision.datasets.MNIST(
    root="data",
    train=True,
    transform=transforms.ToTensor(),
    download=True,
    source_url="https://oneflow-public.oss-cn-beijing.aliyuncs.com/datasets/mnist/MNIST/",
)
mnist_test = flow.utils.vision.datasets.MNIST(
    root="data",
    train=False,
    transform=transforms.ToTensor(),
    download=True,
    source_url="https://oneflow-public.oss-cn-beijing.aliyuncs.com/datasets/mnist/MNIST/",
)

Out:

Downloading https://oneflow-public.oss-cn-beijing.aliyuncs.com/datasets/mnist/MNIST/train-images-idx3-ubyte.gz
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
9913344it [00:00, 36066177.85it/s]
Extracting data/MNIST/raw/train-images-idx3-ubyte.gz to data/MNIST/raw
...

The data will be downloaded and extracted to./data directory.

The oneflow.utils.data.DataLoader wraps an iterable around the dataset.

train_iter = flow.utils.data.DataLoader(
    mnist_train, BATCH_SIZE, shuffle=True
)
test_iter = flow.utils.data.DataLoader(
    mnist_test, BATCH_SIZE, shuffle=False
)
for x, y in train_iter:
    print("x.shape:", x.shape)
    print("y.shape:", y.shape)
    break

Out:

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

Dataset and Dataloader

Building Networks

To define a neural network in OneFlow, we create a class that inherits from nn.Module. We define the layers of the network in the __init__ function and specify how data will pass through the network in the forward function.

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
            nn.ReLU()
        )
    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits
model = NeuralNetwork()
print(model)

Out:

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

Build Network

Training Models

To train a model, we need a loss function (loss_fn) and an optimizer (optimizer). The loss function is used to evaluate the difference between the prediction of the neural network and the real label. The optimizer adjusts the parameters of the neural network to make the prediction closer to the real label (expected answer). Here, we use oneflow.optim.SGD to be our optimizer. This process is called back propagation.

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

The train function is defined for training. In a single training loop, the model makes forward propagation, calculates loss, and backpropagates to update the model’s parameters.

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

We also define a test function to verify the accuracy of the model:

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

We use the train function to begin the train process for several epochs and use the test function to assess the accuracy of the network at the end of each epoch:

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

Out:

loss: 2.299633
loss: 2.303208
loss: 2.298017
loss: 2.297773
loss: 2.294673
loss: 2.295637
Test Error:
 Accuracy: 22.1%, Avg loss: 2.292105
Epoch 2
-------------------------------
loss: 2.288640
loss: 2.286367
...

Autograd Backpropagation and Optimizer

Saving and Loading Models

Use oneflow.save to save the model. The saved model can be then loaded by oneflow.load to make predictions.

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

Model Load and Save

QQ Group

Any problems encountered during the installation or usage, welcome to join the QQ Group to discuss with OneFlow developers and enthusiasts:

Add QQ group by 331883 or scan the QR code below:

Please activate JavaScript for write a comment in LiveRe