From 1e4782a6e9b7149579762b29d3efb054b5762236 Mon Sep 17 00:00:00 2001 From: David Bau Date: Wed, 23 Mar 2022 20:59:53 -0400 Subject: [PATCH] Add module_ and parameter_ names. --- baukit/nethook.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/baukit/nethook.py b/baukit/nethook.py index 12177ca..24da005 100644 --- a/baukit/nethook.py +++ b/baukit/nethook.py @@ -372,6 +372,20 @@ def get_parameter(model, name): raise LookupError(name) +def module_names(model): + """ + Lists all the module names. + """ + return [n for n, _ in model.named_modules()] + + +def parameter_names(model): + """ + Lists all the parameter names. + """ + return [n for n, _ in model.named_parameters()] + + def replace_module(model, name, new_module): """ Replaces the named module within the given model.