-
Notifications
You must be signed in to change notification settings - Fork 77
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
Easy support for per-module train/eval state #37
Comments
Here is one idea of how this could be implemented: ######### in module.py
class ArgsOverride(Module):
def __init__(self, base_module, **kwargs):
self.base_module = base_module
self.kwargs = kwargs
def vars(self, scope):
# presence of ArgsOverride module won’t affect variable names
return self.base_module.vars(scope)
def __call__(self, *args, **kwargs):
kwargs.update(self.kwargs) # in practice this should only
# kwargs which are presented in
# signature of base_module
self.base_module(*args, **kwargs)
######### in utils.py
def reset_args_override(module):
# removes ArgsOverride from module and all submodules
…
############ in user code
model = Resnet50(nclasses=1000)
…
# Here is example how to set most of the network,
# except few modules into training mode
model.block_1.bn_1 = objax.ArgsOverride(model.block_1.bn_1, training=False)
model.block_2.bn_2 = objax.ArgsOverride(model.block_2.bn_2, training=False)
# reset all args override
objax.utils.reset_args_override(model)
# set override on different sub-module
model.block_3.bn_3 = objax.ArgsOverride(model.block_3.bn_3, training=False) @rwightman do you have any feedback about this one? |
@AlexeyKurakin that could work... as implemented ArgsOverride it's quite generic and could be used for any arg. Is there any other arg passed throught the In the absence of additional functionality here, I was likely going to go a subclassing / alternate module impl route ... basically create FrozenBatchNorm, EvalBatchNorm, EvalDropout style classes and helpers to walk module hierarchy within subset of model and switch class types (and copy state). But then that'd run into the checkpoint compat issues discussed. |
Right now I can't really think of other arguments in addition to Also # without ArgsOverride
predict = objax.Jit(lambda x: model(x, training=False), model.vars())
# with ArgsOverride
predict = objax.Jit(objax.ArgsOverride(model, training=False)) |
@rwightman Eventually we decided to use name |
@AlexeyKurakin thanks for the heads up, looks good |
This is tracking issue to improve support of per-module train/eval state in Objax.
This issue originates from discussion of PyTorch vs Objax-style of propagating training/eval mode #29
PyTorch style allows easy way to specify per-module train/eval mode, like in the following example:
There is no clean way to achieve the same thing in Objax right now. One possibility is to use
functools.partial
:However there are some problems with
functools.partial
:vars()
are not propagated).bn_eval = functools.partial(bn, training=False)
and caller will try to passtraining
argument tobn_eval
it will cause run-time errorfunctools.partial
after it applied to a moduleThus we need a better solution to do per-module train/eval state
The text was updated successfully, but these errors were encountered: