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.

Sign up for GitHub

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

Jump to bottom

[WIP]Use get_name in forward hook #393

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

[WIP]Use get_name in forward hook #393

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9e6a581
get_name
NihalHarish Oct 30, 2020
3c9b1ad
resolve import
NihalHarish Oct 30, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 11 additions & 2 deletions smdebug/pytorch/hook.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
# Third Party
import torch
import torch.distributed as dist
from torch.nn.parallel.data_parallel import DataParallel

# First Party
from smdebug.core.collection import DEFAULT_PYTORCH_COLLECTIONS, CollectionKeys
Expand Down Expand Up @@ -197,6 +198,14 @@ def register_hook(self, module):
# for compatibility with ZCC patches which call this
self.register_module(module)

@staticmethod
def _add_module_name(module, module_name):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

any reason why static?

if isinstance(module, DataParallel):
module.module._module_name = module_name
else:
module._module_name = module_name
return module

def register_module(self, module):
"""
This function registers the forward hook. If user wants to register the hook
Expand All @@ -215,9 +224,9 @@ def register_module(self, module):

for name, submodule in module.named_modules():
assert submodule not in self.module_set, f"Don't register module={module} twice"
submodule._module_name = name
Hook._add_module_name(submodule, name)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

self. ?

self.module_set.add(submodule)
module._module_name = module._get_name()
Hook._add_module_name(module, module._get_name())
self.module_set.add(module)

# Use `forward_pre_hook` for the entire net
Expand Down