Easy support for per-module train/eval state about objax HOT 5 CLOSED

######### 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)

# without ArgsOverride
predict = objax.Jit(lambda x: model(x, training=False), model.vars())

# with ArgsOverride
predict = objax.Jit(objax.ArgsOverride(model, training=False))