OneFlow is a deep learning framework engineered for large-scale distributed training. It employs a static computation graph and provides efficient automatic differentiation. A core principle of OneFlow is "write once, run anywhere," enabling model code to execute seamlessly across different hardware and distributed setups without modification.

Key features of OneFlow include:

• Performance: Optimized computation graphs and memory management deliver high efficiency for both training and inference.
• Flexibility: The framework supports custom operators and model architectures, facilitating complex experimentation.
• Distributed Training: Built-in strategies simplify scaling model training to large computing clusters.
• Usability: A clean API and comprehensive documentation lower the learning barrier for users.

Implementing an Image Classifier with OneFlow

This example demonstrates how to build a simple image classifier for the CIFAR-10 dataset.

Begin by installing the framework:

pip install oneflow

Next, import the required modules and define the neural network model:

import oneflow as flow
import oneflow.nn as nn
import oneflow.optim as optim
from oneflow.utils.vision import datasets, transforms

class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(3, 6, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(6, 16, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(16 * 5 * 5, 120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, 10)
        )

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

net = ConvNet()
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

Prepare the CIFAR-10 dataset with data transformations:

# Define data transformations
transform_pipeline = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# Load datasets
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_pipeline)
train_loader = flow.utils.data.DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=2)

test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_pipeline)
test_loader = flow.utils.data.DataLoader(test_dataset, batch_size=4, shuffle=False, num_workers=2)

Train the model for a specified number of epochs:

num_epochs = 2
for epoch in range(num_epochs):
    total_loss = 0.0
    for batch_idx, (images, targets) in enumerate(train_loader):
        # Clear gradients from the previous step
        optimizer.zero_grad()
        
        # Forward pass
        predictions = net(images)
        loss = loss_fn(predictions, targets)
        
        # Backward pass and optimization
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
        if batch_idx % 2000 == 1999:  # Print every 2000 mini-batches
            print(f'Epoch [{epoch+1}], Batch [{batch_idx+1}], Loss: {total_loss / 2000:.3f}')
            total_loss = 0.0