oneflow.nn.graph.graph_block.GraphModule.activation_checkpointing¶
-
property
GraphModule.
activation_checkpointing
¶ Set/Get whether do activation checkpointing in this nn.Module.
For example:
import oneflow as flow from oneflow.nn.graph import GraphModule class Graph(flow.nn.Graph): def __init__(self): super().__init__() self.linear1 = flow.nn.Linear(3, 5, False) self.linear2 = flow.nn.Linear(5, 8, False) self.linear1.to(GraphModule).activation_checkpointing = True self.linear2.to(GraphModule).activation_checkpointing = True def build(self, x): y_pred = self.linear1(x) y_pred = self.linear2(y_pred) return y_pred graph = Graph()