torch.nn.register_module_module_registration_hook
function register_module_module_registration_hook(hook: ModuleRegistrationHook): RemovableHandleRegisters a global hook that executes whenever a submodule is added to any module.
Installs a hook that will be called for ALL modules when add_module() is invoked (i.e., when assigning a submodule as a property). This hook can inspect, wrap, or modify submodules before they're registered. Can also prevent registration by returning null. Useful for:
- Architecture monitoring: Tracking module hierarchy as it's built
- Module wrapping: Adding logging, profiling, or other wrappers automatically
- Module validation: Ensuring submodules meet compatibility requirements
- Dynamic architecture: Modifying/replacing modules during construction
- Debugging: Logging the model structure as it's created
The hook receives the parent module, the property name (e.g., 'layer1'), and the submodule being registered. Returning null prevents registration. Returning void registers unchanged. Returning a module registers the modified/wrapped module instead.
- Module construction time: Called when building the model structure
- Can modify/replace: Can return modified module or wrap with decorators
- Parent context: Has access to parent module to understand hierarchy
- Property naming: Name parameter is the assigned property name
- Breaking model structure: Careful with module modifications
- Type compatibility: Wrapped modules must be compatible with parent expectations
- Initialization order: Runs during model construction, affects initial setup
Parameters
hookModuleRegistrationHook- ModuleRegistrationHook called with (module, name, submodule) when add_module() is called
Returns
RemovableHandle– RemovableHandle to unregister this hook using .remove()Examples
// Log all submodule registration
const hook = (module, name, submodule) => {
if (submodule) {
console.log(`Adding ${submodule.constructor.name} as '${name}' to ${module.constructor.name}`);
}
};
torch.nn.register_module_module_registration_hook(hook);// Wrap modules with logging
const hook = (module, name, submodule) => {
if (submodule && submodule instanceof torch.nn.Linear) {
// Return wrapped module
return new LoggingWrapper(submodule);
}
};
torch.nn.register_module_module_registration_hook(hook);// Monitor model growth
let module_count = 0;
const hook = (module, name, submodule) => {
if (submodule) {
module_count++;
console.log(`Total modules: ${module_count}`);
}
};
torch.nn.register_module_module_registration_hook(hook);See Also
- PyTorch torch.nn.modules.module.register_module_module_registration_hook
- register_module_buffer_registration_hook - Hook for buffers
- register_module_parameter_registration_hook - Hook for parameters
- Module.add_module - Method that triggers this hook