Mocking PyTorch with OneFlow¶
With the alignment of OneFlow API and PyTorch, users can easily migrate the model from PyTorch to OneFlow. In this article, we will introduce three methods to port PyTorch code to OneFlow.
1. import oneflow as torch¶
Change the original code
import torch to:
import oneflow as torch
You can train the original model with OneFlow; however, you have to manually modify all the files that contain
import torch by this approach. In addition, it is necessary to modify the source code if a third-party library uses
2. using the command-line tool¶
Oneflow provides a command-line tool, which helps to mock the environment of PyTorch within OneFlow’s Python Package and forward references of Pytorch to the real OneFlow module. The specific steps are as follows:
Enabling the mocking of PyTorch
eval $(python3 -m oneflow.mock_torch)
To facilitate debugging, OneFlow provides two parameters to this method:
lazy parameter. When
lazy=True, a mock object will be returned without an immediate error for non-existent interfaces.
verbose parameter. If
verbose=Trueis set simultaneously, it will print out which mock objects are accessed or used for debugging.
The usage is as follows:
Enabling the mocking of PyTorch, and configuring lazy and verbose parameters.
eval $(oneflow-mock-torch --lazy --verbose)
eval $(python3 -m oneflow.mock_torch --lazy --verbose)
After running the above command, you can observe the effect in the following example.
import torch print(torch.__file__) import oneflow as flow x = torch.zeros(2, 3) print(isinstance(x, flow.Tensor))
Disabling the mocking of PyTorch
eval $(oneflow-mock-torch disable)
eval $(python3 -m oneflow.mock_torch disable)
3. using built-in functions of OneFlow¶
We provide the mock function with fine granularity, and the users can determine whether to enable this function for a piece of code.
In the following
with statement, the PyTorch module imported is OneFlow.
import oneflow.mock_torch as mock with mock.enable(): import torch print(torch.__file__) import oneflow as flow x = torch.zeros(2, 3) print(isinstance(x, flow.Tensor))
Similarly, OneFlow provides parameters lazy and verbose for
mock.enable() to facilitate debugging, which can be set like this:
with mock.enable(lazy=True, verbose=True)
You can turn off the mock function like this when it’s needed to use the real torch module.
with mock.disable(): import torch print(torch.__file__)
mock.disable can act as functions. For example, if you want to train a model with OneFlow, but it needs to be loaded by PyTorch. Then, you can use it with the following code:
mock.enable() ... with mock.disable() module = torch.load_module(...) # train the module with oneflow
A dictionary with the value of module is saved separately in enable and disable mode. When you turn enable or disable on and off, the dictionary will replace
sys.modules and the global variables that the current module belongs to. Therefore, users are required to
import the module they need in each mode, and the code below will raise an error
name 'torch' is not defined in
with statment of disable.
with mock.enable(): import torch with mock.disable(): torch.ones(2, 3)
With the alignment of OneFlow API and PyTorch, users can easily migrate the PyTorch code to OneFlow. As mentioned above, three methods are available to train the PyTorch model with OneFlow. This is how users can experience the ultimate performance of OneFlow.