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 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.
oneflow.nn.Graph
Eager 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
methodInstantiate 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 it¶
Initializes internal Graph states. |
|
The |
|
Add an optimizer, an learning rate scheduler to the graph. |
|
Set the GradScaler for gradient and loss scaling. |
Execute Graph¶
Call nn.Graph subclass instance to run your customized graph. |
Graph Config option¶
If set to true, then graph will use mixed precision mode, it means use both float16 and float32 during model training. |
|
Enable ZeRO redundancy optimizer. |
|
If set to true, try to fuse cast + scale + l1_l2_regularize_gradient + model_update to one op to improve performance. |
|
If set to true, try to fuse a binary element-wise add operetor to one of the predecessors to improve performance. |
|
If set to true, try to fuse cast and scalar_mul_by_tensor to improve performance. |
|
Set num of steps to accumulate gradient. |
|
Whether enable cudnn conv operatioin to use heuristic search algorithm. |
Block Config option¶
Set stage id and placement of nn.Module/ModuleBlock in pipeline parallelism. |
|
Set/Get whether do activation checkpointing in this nn.Module. |
Save & Load Model¶
Returns a dictionary containing a whole state of the graph. |
|
Copies module’s states and other graph states from |