| |
|
| |
|
| |
|
| |
|
| |
|
| | from caffe2.python import scope |
| |
|
| | import contextlib |
| | import logging |
| |
|
| | logger = logging.getLogger(__name__) |
| |
|
| |
|
| | class ParameterSharingContext(object): |
| | """ |
| | This class manages scope driven way of parameter sharing across different |
| | NameScopes. |
| | """ |
| |
|
| | def __init__(self): |
| | self._scope_overrides = {} |
| | self._contexts = [] |
| |
|
| | def _resolve_scope_overrides(self, candidate_scope): |
| | """ |
| | Recursively resolves all scope overrides, i.e multiple steps of |
| | override can be used. |
| | |
| | For example, if one provides following scope overrides: |
| | {'scope_b': 'scope_a'} and within 'scope_b' - {'shared_child': ''}, |
| | then name 'w' will get resolved to the following blobs depending on the |
| | namescope: |
| | a. 'scope_a' -> 'scope_a/w' |
| | b. 'scope_b' -> 'scope_a/w' |
| | c. 'scope_c' -> 'scope_c/w' |
| | d. 'scope_b/shared_child' -> 'scope_a/w' |
| | d. 'scope_b/unshared_child' -> 'scope_a/unshared_child/w' |
| | """ |
| | best_scope = candidate_scope |
| | best_scope_idx = 0 |
| | sub_scopes = candidate_scope.split(scope._NAMESCOPE_SEPARATOR) |
| |
|
| | cur_scope = '' |
| | for idx, sub_scope in enumerate(sub_scopes): |
| | cur_scope = cur_scope + sub_scope + scope._NAMESCOPE_SEPARATOR |
| | if cur_scope in self._scope_overrides: |
| | best_scope = self._scope_overrides[cur_scope] |
| | best_scope_idx = idx |
| | if best_scope == candidate_scope: |
| | return candidate_scope |
| | else: |
| | return (self._resolve_scope_overrides(best_scope) + |
| | scope._NAMESCOPE_SEPARATOR.join( |
| | sub_scopes[best_scope_idx + 1:])) |
| |
|
| | def get_parameter_name(self, name): |
| | candidate_scope = scope.CurrentNameScope() |
| | best_scope = self._resolve_scope_overrides(candidate_scope) |
| | if best_scope != candidate_scope: |
| | logger.info("Overwriting scope {0} with scope {1}".format( |
| | candidate_scope, best_scope)) |
| |
|
| | return best_scope + name |
| |
|
| | def add_scope_overrides(self, shared_scopes): |
| | self._contexts.append(shared_scopes) |
| | self._scope_overrides.update(shared_scopes) |
| |
|
| | def pop(self): |
| | assert len(self._contexts) > 0 |
| | self._contexts.pop() |
| | self._scope_overrides = {} |
| | for x in self._contexts: |
| | self._scope_overrides.update(x) |
| |
|
| |
|
| | parameter_sharing_context = ParameterSharingContext() |
| |
|
| |
|
| | def _normalize_namescope(namescope): |
| | if namescope and namescope[-1] != scope._NAMESCOPE_SEPARATOR: |
| | return namescope + scope._NAMESCOPE_SEPARATOR |
| | else: |
| | return namescope |
| |
|
| |
|
| | @contextlib.contextmanager |
| | def ParameterSharing(shared_scopes): |
| | """ |
| | Helper function for sharing scopes. |
| | All the parameters within the shared_scopes, will be remapped with the |
| | respect of CurrentNamescope() |
| | |
| | I.e. if one calls ParameterSharing with {'scope_b': 'scope_'a'}, from the |
| | scope 'some_global_scope', it'll effectively mean, that all parameters from |
| | 'some_global_scope/scope_b' will shared with the parameters from |
| | 'some_global_scope/scope_a' |
| | """ |
| | assert isinstance(shared_scopes, dict) |
| |
|
| | shared_scope_overrides = {} |
| | current_scope = scope.CurrentNameScope() |
| | for k, v in shared_scopes.items(): |
| | assert not v.startswith(k), ( |
| | "Illegal override for parameter sharing. {} is prefix of {}". |
| | format(k, v)) |
| | k = current_scope + k |
| | v = current_scope + v |
| | |
| | k = _normalize_namescope(k) |
| | v = _normalize_namescope(v) |
| | shared_scope_overrides[k] = v |
| |
|
| | try: |
| | parameter_sharing_context.add_scope_overrides(shared_scope_overrides) |
| | yield |
| | finally: |
| | parameter_sharing_context.pop() |
| |
|
