Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

use object.__setattr__ in DualModule and DeployableModule to avoid unexpected key in self._modules #468

Merged
merged 6 commits into from
Dec 27, 2023

Conversation

marigoold
Copy link
Contributor

@marigoold marigoold commented Dec 27, 2023

之前直接赋值 self.xxx = some_nn_Module 的写法,会在 self._modules 中写入一个 key='xxx' 的键值对,这是由于 torch.nn.Module 的 __setattr__ 决定的,导致遍历、取 state_dict 等行为得到的结果与期望值不符,会多一层不必要的 _torch_module 或者 _deployable_module_model,比如下列场景:

In [1]: from onediff.infer_compiler import oneflow_compile

In [2]: from torchvision.models.resnet import resnet50

In [3]: model = resnet50(pretrained=False, progress=True).eval()

In [4]: of_model = oneflow_compile(model, use_graph=True)

In [5]: of_model._modules.keys()
Out[5]: odict_keys(['_deployable_module_model'])

In [6]: of_model._deployable_module_model._modules.keys()
Out[6]: odict_keys(['_torch_module'])                                                                                                                                                                                                                       

修改之后不用直接赋值语句,而是调用 object.__setattr__,避免在 self._modules 中增加不必要的键值对,然后把子 module 中的 _modules、_parameters、_buffers 都浅拷贝一份,这样遍历、state_dict、get_submodule 等行为都正常了

@marigoold marigoold merged commit 2d7793f into main Dec 27, 2023
0 of 2 checks passed
@marigoold marigoold deleted the dev_wy_copy_modules_in_torch_module branch December 27, 2023 12:47
hjchen2 added a commit that referenced this pull request Dec 28, 2023
…avoid unexpected key in self._modules (#468)"

This reverts commit 2d7793f.
marigoold added a commit that referenced this pull request Dec 29, 2023
…avoid unexpected key in self._modules (#468)"

This reverts commit 2d7793f.
marigoold added a commit that referenced this pull request Dec 29, 2023
…avoid unexpected key in self._modules" (#476)

Reverts #468
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants