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

Graph freezing and unfreezing #614

Merged
merged 106 commits into from
May 7, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
106 commits
Select commit Hold shift + click to select a range
de587bd
Manual rebase of dev-module-composition on top of current master
tkornuta-nvidia Mar 11, 2020
15b3b11
Adding neural graph examples
tkornuta-nvidia Mar 24, 2020
2756a66
formatting fixes
tkornuta-nvidia Mar 24, 2020
5bad9c9
new proposal
blisc Apr 1, 2020
e26a906
Merge branch 'master' of github.com:NVIDIA/NeMo into fix-sign
tkornuta-nvidia Apr 6, 2020
85bb729
Merge branch 'fix-sign' of github.com:NVIDIA/NeMo into fix-sign
tkornuta-nvidia Apr 6, 2020
0576ce7
format fix of Jason's example
tkornuta-nvidia Apr 6, 2020
3f1ac01
Added modules recording, with retrieval of module name
tkornuta-nvidia Apr 7, 2020
a57b34d
keyerror comment
tkornuta-nvidia Apr 7, 2020
fc3deac
Introduced name property of modules
tkornuta-nvidia Apr 7, 2020
5221fef
implemented NameRegistry basic functionality, test5 working, unit tes…
tkornuta-nvidia Apr 7, 2020
dd9af75
example 5 working, ObjectRegistry finished, app_state moved to utils
tkornuta-nvidia Apr 8, 2020
ab35621
test app_state fixed
tkornuta-nvidia Apr 8, 2020
bbd3589
TokenClassification self.name fix
tkornuta-nvidia Apr 8, 2020
56efff5
ObjectRegistry tests
tkornuta-nvidia Apr 8, 2020
c838f13
cleanup of neural graphs class, NG manager inheriting from ObjectRegi…
tkornuta-nvidia Apr 8, 2020
1153001
fixes to graph and manager, NeMo tests passing (locally)
tkornuta-nvidia Apr 8, 2020
9df6568
code formatting
tkornuta-nvidia Apr 8, 2020
868ec8a
turning examples into unit/integration tests, small fixes here and th…
tkornuta-nvidia Apr 9, 2020
51e852c
Graph and module nesting - operation mode injected, implemented funct…
tkornuta-nvidia Apr 9, 2020
e2ed33a
formatting fix
tkornuta-nvidia Apr 9, 2020
069dd10
merged with master
tkornuta-nvidia Apr 9, 2020
a054311
cleanup of output port logic in __call__, other cleanups, LGTM fixes
tkornuta-nvidia Apr 9, 2020
c0e5f0b
Further cleanup of neural_modules.py
tkornuta-nvidia Apr 9, 2020
98cfe43
LGTM cleanups
tkornuta-nvidia Apr 9, 2020
d109b5c
style fix:]
tkornuta-nvidia Apr 9, 2020
6487a65
micro cleanup
tkornuta-nvidia Apr 9, 2020
6c16c6e
micro cleanup2
tkornuta-nvidia Apr 9, 2020
b7d3f1c
punctuation module name fix
tkornuta-nvidia Apr 9, 2020
5548495
Working on NmTensor extensions
tkornuta-nvidia Apr 10, 2020
c835279
style fix
tkornuta-nvidia Apr 10, 2020
70320fe
Cleaned up the call output tensor logic in neural module and graph cl…
tkornuta-nvidia Apr 10, 2020
c9b2585
NmTensor test cleanup
tkornuta-nvidia Apr 10, 2020
adeeccc
style fix
tkornuta-nvidia Apr 10, 2020
127307b
LGTM cleanups
tkornuta-nvidia Apr 10, 2020
7b52fbe
style fixes
tkornuta-nvidia Apr 10, 2020
81b87ea
BoundOutputs class and unit tests
tkornuta-nvidia Apr 10, 2020
fbfc87e
style fixes
tkornuta-nvidia Apr 10, 2020
41ee935
manual output port binding operational
tkornuta-nvidia Apr 10, 2020
bce7e90
manual output port binding operational
tkornuta-nvidia Apr 10, 2020
752b810
Recording and updating all modules during connecting to a graph
tkornuta-nvidia Apr 14, 2020
f68c883
work in progress on output port binding, not operations
tkornuta-nvidia Apr 15, 2020
9c3f78b
bound output test fix
tkornuta-nvidia Apr 15, 2020
b8bf040
preparing ground for tensor copy
tkornuta-nvidia Apr 15, 2020
691e90b
Rewritten graph nesting - works by executing inner graph modules' cal…
tkornuta-nvidia Apr 18, 2020
ec2bfe6
fixed tests, formatted the code
tkornuta-nvidia Apr 18, 2020
cb250da
refactoring and cleanup of neural graphs
tkornuta-nvidia Apr 21, 2020
6907cd7
Refactoring, cleanups of NeuralGraphs nesting/binding
tkornuta-nvidia Apr 21, 2020
cb69367
reformating fix
tkornuta-nvidia Apr 21, 2020
8389f21
Fixes, unit tests for binding of nested graphs, automatic and manual,…
tkornuta-nvidia Apr 22, 2020
9235ca8
Added several tests for graph nesting, covered (and fixed) case of pa…
tkornuta-nvidia Apr 22, 2020
59b1d75
final tests for graph nesting
tkornuta-nvidia Apr 22, 2020
fc9d995
style fix
tkornuta-nvidia Apr 22, 2020
ece11c3
refactored the neural module serialization, made it more modular, so …
tkornuta-nvidia Apr 23, 2020
cd703de
style fix
tkornuta-nvidia Apr 23, 2020
465979d
graph serialization 80%
tkornuta-nvidia Apr 23, 2020
ba46466
graph serialization operational
tkornuta-nvidia Apr 23, 2020
fa1e4b3
work on deserialization, working up to the execution of graph
tkornuta-nvidia Apr 24, 2020
b810243
format fix
tkornuta-nvidia Apr 24, 2020
9f46e20
serialization and deserialization of JASPER (processor, encoder and d…
tkornuta-nvidia Apr 27, 2020
77ad78d
Merge remote-tracking branch 'origin/master' into merge
tkornuta-nvidia Apr 27, 2020
7b5b9d4
merged with master and reformatted
tkornuta-nvidia Apr 27, 2020
26a7621
Updated EN and ZH:] version of documentation for configuration/custom…
tkornuta-nvidia Apr 28, 2020
314cf10
Unification of singleton meta class
tkornuta-nvidia Apr 28, 2020
7f0ee86
Added to the train() function signatures
tkornuta-nvidia Apr 28, 2020
4076808
formatting fix
tkornuta-nvidia Apr 28, 2020
41bb68f
LGTM fixes
tkornuta-nvidia Apr 28, 2020
cb56b35
line numbers
tkornuta-nvidia Apr 28, 2020
5ff3896
Tweak of deserialize - moved name and overwrite_params
tkornuta-nvidia Apr 28, 2020
c955a7a
style fix
tkornuta-nvidia Apr 28, 2020
3b47bd9
Commeting the condition for both None: eval() seems to be calling :]
tkornuta-nvidia Apr 28, 2020
fd88c70
LGTM fix
tkornuta-nvidia Apr 28, 2020
98b03cf
method comment
tkornuta-nvidia Apr 28, 2020
18077fc
reorganization of directories in unit/core/utils
tkornuta-nvidia Apr 28, 2020
a7a04e3
Fix of NM init_params collection, made it much more robust and faster…
tkornuta-nvidia Apr 28, 2020
a932c8a
Renamed NG integration tests
tkornuta-nvidia Apr 28, 2020
7c57cdf
formatting
tkornuta-nvidia Apr 28, 2020
6d785f1
additional tests
tkornuta-nvidia Apr 28, 2020
5633963
style fix
tkornuta-nvidia Apr 28, 2020
8f6ddc5
A simple test for graph config import and export
tkornuta-nvidia Apr 28, 2020
1977db0
import/export/serialization/deserialization tests cnt, minor fixes he…
tkornuta-nvidia Apr 29, 2020
bd6c12e
minor cleanup, unit test: a simple graph with a loop (not passing)
tkornuta-nvidia Apr 29, 2020
7463dd8
style fix
tkornuta-nvidia Apr 29, 2020
17e4dd9
Refactored the whole solution, switching from module_name:port_name t…
tkornuta-nvidia Apr 29, 2020
7b94c57
format fix
tkornuta-nvidia Apr 29, 2020
e5917fe
Added neural type export to config file for both graph bound input an…
tkornuta-nvidia Apr 29, 2020
35b31c9
jasper polished
tkornuta-nvidia Apr 29, 2020
5375055
PR polish, added type hints to all classes, added graph summary
tkornuta-nvidia May 1, 2020
9f6ee98
removed 'name' in for in actions.py
tkornuta-nvidia May 1, 2020
45dacd3
NG-related integration tests
tkornuta-nvidia May 1, 2020
356e0c2
LGTM fix
tkornuta-nvidia May 1, 2020
cd1e23e
bind description updated
tkornuta-nvidia May 1, 2020
37ba923
removed app_state from jasper
tkornuta-nvidia May 1, 2020
e28979e
graph input_ports and output_ports are now immutable, added some unit…
tkornuta-nvidia May 1, 2020
693b744
removed a line - formatting fix :]
tkornuta-nvidia May 2, 2020
413c0fa
name fix in test
tkornuta-nvidia May 2, 2020
ccc1b21
Minor touches here and there
tkornuta-nvidia May 2, 2020
a191dbe
Unique names generates from names of classes
tkornuta-nvidia May 2, 2020
68bd7e2
moved graph managed to utils, removed it from init
tkornuta-nvidia May 4, 2020
bc216db
Moved neural_graph_manager to utils that enable to remove it (along w…
tkornuta-nvidia May 4, 2020
50bc6f0
merged with graph master
tkornuta-nvidia May 5, 2020
e170915
added freeze and unfreeze to NeuralGraphs
tkornuta-nvidia May 6, 2020
8b728b7
removed the examples
tkornuta-nvidia May 6, 2020
a4ac16b
formatting fix
tkornuta-nvidia May 6, 2020
c012c85
reverted changes from improper merge
tkornuta-nvidia May 6, 2020
72a5972
added nontrainable
tkornuta-nvidia May 7, 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
15 changes: 14 additions & 1 deletion nemo/backends/pytorch/nm.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import torch as t
import torch.nn as nn

from nemo.core import DeviceType, NeuralModule, WeightShareTransform
from nemo.core import DeviceType, ModuleType, NeuralModule, WeightShareTransform
okuchaiev marked this conversation as resolved.
Show resolved Hide resolved
from nemo.utils.helpers import get_cuda_device, rgetattr, rsetattr


Expand Down Expand Up @@ -38,6 +38,9 @@ def __init__(self, pretrained_model_name=None, name=None):
NeuralModule.__init__(self, name) # For NeuralModule API
nn.Module.__init__(self) # For PyTorch API

# Set module type.
self._type = ModuleType.trainable

self._device = get_cuda_device(self.placement)

# Store pretrained model name (to be removed/changed)
Expand Down Expand Up @@ -132,6 +135,8 @@ class NonTrainableNM(NeuralModule):
def __init__(self, name=None):
NeuralModule.__init__(self, name) # For NeuralModule API
self._device = get_cuda_device(self.placement)
# Set module type.
self._type = ModuleType.nontrainable

def __call__(self, force_pt=False, *input, **kwargs):
pt_call = len(input) > 0 or force_pt
Expand Down Expand Up @@ -191,6 +196,10 @@ class DataLayerNM(NeuralModule):

def __init__(self, name=None):
NeuralModule.__init__(self, name) # For NeuralModule API

# Set module type.
self._type = ModuleType.datalayer

self._device = get_cuda_device(self.placement)

# if 'batch_size' not in kwargs:
Expand Down Expand Up @@ -326,6 +335,10 @@ class LossNM(NeuralModule):

def __init__(self, name=None):
NeuralModule.__init__(self, name) # For NeuralModule API

# Set module type.
self._type = ModuleType.loss

self._device = get_cuda_device(self.placement)

def get_weights(self):
Expand Down
44 changes: 43 additions & 1 deletion nemo/core/neural_graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@

from nemo.core import OperationMode
from nemo.core.neural_interface import NeuralInterface
from nemo.core.neural_modules import NeuralModule
from nemo.core.neural_modules import ModuleType, NeuralModule
from nemo.core.neural_types import NeuralPortNameMismatchError, NeuralType, NmTensor
from nemo.package_info import __version__ as nemo_version
from nemo.utils import logging
Expand Down Expand Up @@ -920,3 +920,45 @@ def summary(self) -> str:

# Return the result.
return desc

def freeze(self, module_names: Optional[List[str]] = None):
"""
A method that freezes the weights of the trainable modules in a graph.
Args:
module_names: List of modules to be frozen (Optional). If passed, all modules will be unfrozen.
Raises:
KeyError: If name of the module won't be recognized.
"""
if module_names is None:
# Work on all modules.
module_names = self._modules.keys()
# Iterate through modules one by one.
for name in module_names:
if name not in self._modules.keys():
raise KeyError("Module `{}` not present in the `{}` graph".format(name, self.name))
# Check module type.
module = self._modules[name]
if module.type == ModuleType.trainable:
# Freeze weights of the module
module.freeze()

def unfreeze(self, module_names: Optional[List[str]] = None):
"""
Unfreezes weights of the trainable modules in a graph.
Args:
module_names: List of modules to be unfrozen (Optional). If not passed, all modules will be unfrozen.
Raises:
KeyError: If name of the module won't be recognized.
"""
if module_names is None:
# Work on all modules.
module_names = self._modules.keys()
# Iterate through modules one by one.
for name in module_names:
if name not in self._modules.keys():
raise KeyError("Module `{}` not present in the `{}` graph".format(name, self.name))
# Check module type.
module = self._modules[name]
if module.type == ModuleType.trainable:
# Unfreeze weights of the module.
module.unfreeze()
30 changes: 19 additions & 11 deletions nemo/core/neural_modules.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.

"""This file contains NeuralModule and NmTensor classes."""
__all__ = ['WeightShareTransform', 'NeuralModule']
__all__ = ['WeightShareTransform', 'NeuralModule', 'ModuleType']

import uuid
from abc import abstractmethod
Expand All @@ -39,6 +38,16 @@
YAML = YAML(typ='safe')


class ModuleType(Enum):
""" Back-end independent module types """

module = 0
okuchaiev marked this conversation as resolved.
Show resolved Hide resolved
datalayer = 1
trainable = 2
loss = 3
nontrainable = 4


class WeightShareTransform(Enum):
"""When sharing parameters, what kind of transform to apply."""

Expand Down Expand Up @@ -69,6 +78,9 @@ def __init__(self, name=None):
# Register module and store the generated name.
self._name = self._app_state.register_module(self, name)

# Set "module" type as default.
self._type = ModuleType.module

# Set "both" as default operation mode.
self._operation_mode = OperationMode.both

Expand Down Expand Up @@ -478,15 +490,6 @@ def _deserialize_configuration(cls, serialized_init_params: Dict[str, Any]):
# In this case configuration = init parameters.
return serialized_init_params

@deprecated(version=0.11)
@staticmethod
def create_ports(**kwargs):
""" Deprecated method, to be remoted in the next release."""
raise Exception(
'Deprecated method. Please implement ``inputs`` and ``outputs`` \
properties to define module ports instead'
)

@property
@abstractmethod
def input_ports(self) -> Dict[str, NeuralType]:
Expand Down Expand Up @@ -534,6 +537,11 @@ def operation_mode(self):
""" Returns the operation mode. """
return self._operation_mode

@property
def type(self):
""" Returns the type of module. """
return self._type

@operation_mode.setter
def operation_mode(self, operation_mode: OperationMode):
""" Sets the operation mode. """
Expand Down