Skip to content

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

Merged
marigoold merged 6 commits intomainfrom
dev_wy_copy_modules_in_torch_module
Dec 27, 2023
Merged

use object.__setattr__ in DualModule and DeployableModule to avoid unexpected key in self._modules#468
marigoold merged 6 commits intomainfrom
dev_wy_copy_modules_in_torch_module

Conversation

@marigoold
Copy link
Collaborator

@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
@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