oneflow.nn.graph.graph_config.GraphConfig.set_gradient_accumulation_steps

GraphConfig.set_gradient_accumulation_steps(value)

Set num of steps to accumulate gradient.

For example:

import oneflow as flow

class Graph(flow.nn.Graph):
    def __init__(self):
        super().__init__()
        self.linear = flow.nn.Linear(3, 8, False)
        # Let graph do gradient accumulation, such as pipelining parallelism depends on gradient accumulation.
        self.config.set_gradient_accumulation_steps(4)
    def build(self, x):
        return self.linear(x)

graph = Graph()

Parameters

value (int) – num of steps.