oneflow.nn.Graph

Base class for running neural networks in Static Graph Mode.

Currently, there are two main ways to run models in deep learning frameworks, namely dynamic graphs and static graphs , which are also conventionally referred to as Eager Mode to Static Graph Mode and Static Graph Mode in OneFlow.

Both approaches have their advantages and disadvantages, and OneFlow provides support for both approaches, with Eager mode being the default.

Generally speaking, dynamic graphs are easier to use and static graphs have more performance advantages. oneflow.nn.Graph module is provided by OneFlow to allow users to build static graphs and train models with Eager-like programming conventions.

Eager Mode to Static Graph Mode

OneFlow runs in Eager mode by default.

OneFlow’s nn.Graph is programmed in a style very similar to Eager Mode, so it is possible to make small changes and get large performance gains.

The following script shows the process of building a neural network in eager mode using the interface under oneflow.nn :

import oneflow as flow
import oneflow.nn as nn

class ModuleMyLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.weight = nn.Parameter(flow.randn(in_features, out_features))
        self.bias = nn.Parameter(flow.randn(out_features))

    def forward(self, input):
        return flow.matmul(input, self.weight) + self.bias

linear_model = ModuleMyLinear(4, 3)

Eager nn.Module can be reused by nn.Graph. The above script for eager mode can be changed to static Graph mode by adding just a few lines of code, which consists of the following steps:

  • Define your customized graph as a subclass of nn.Graph

  • At the beginning of __init__. Call super().__init__() to let OneFlow do the necessary initialization of the Graph

  • Reuse the nn.Module object in Eager mode in __init__ (self.model = model)

  • Describe the computation in the build method

  • Instantiate your graph then call it.

class GraphMyLinear(nn.Graph):
    def __init__(self):
        super().__init__()
        self.model = linear_model

    def build(self, input):
        return self.model(input)

graph_mylinear = GraphMyLinear()
input = flow.randn(1, 4)
out = graph_mylinear(input)
print(out)

tensor([[-0.3298, -3.7907,  0.1661]], dtype=oneflow.float32)

Static Graph Mode

Constructing a Graph

Base class for training or evaluating a neural network in static graph mode.

__init__

Initializes internal Graph states.

build

The build() method must be overridden to define neural network computaion logic.

add_optimizer

Add an optimizer, an learning rate scheduler to the graph.

set_grad_scaler

Set the GradScaler for gradient and loss scaling.

Executing a Graph

Call a nn.Graph instance to run a customized graph.

__call__

Call nn.Graph subclass instance to run your customized graph.

Config options on a Graph

Optimization options of a nn.Graph.

enable_amp

If set to true, then graph will use mixed precision mode, it means use both float16 and float32 during model training.

enable_zero

Enable ZeRO redundancy optimizer.

allow_fuse_model_update_ops

If set to true, try to fuse cast + scale + l1_l2_regularize_gradient + model_update to one op to improve performance.

allow_fuse_add_to_output

If set to true, try to fuse a binary element-wise add operator to one of the predecessors to improve performance.

allow_fuse_cast_scale

If set to true, try to fuse cast and scalar_mul_by_tensor to improve performance.

set_gradient_accumulation_steps

Set num of steps to accumulate gradient.

enable_cudnn_conv_heuristic_search_algo

Whether enable cudnn conv operation to use heuristic search algorithm.

enable_straighten_algorithm

Whether enable the straighten algorithm.

enable_compress_memory

If true, then the graph will try its best to find the minimum memory allocation strategy.

Config options on a GraphModule

GraphModule is the graph representation of a nn.Module in a nn.Graph.

When an nn.Module is added into an nn.Graph, it is wrapped into a ProxyModule. The ProxyModule has a GraphModule inside it. You can get and set the GraphModule to enable graph optimization on the nn.Module.

set_stage

Set stage id and placement of nn.Module in pipeline parallelism.

activation_checkpointing

Set/Get whether do activation checkpointing in this nn.Module.

Save & Load a Model

state_dict

Returns a dictionary containing a whole state of the graph.

load_state_dict

Copies module’s states and other graph states from state_dict into this graph.

Debug a Graph

__repr__

For printing the graph structure.

debug

Open or close debug mode of the graph.

name

Name auto-generated for this graph.