A .py utility file for neural network learing rate policies I've created a .py utility file, which specifies the learning rate policy for the neural network in PyTorch. The program in prior reads a json and creates a dict, called here is cfg.
So far, I have implemented 3 kinds of learning policies - steps, exp, cos, and for each of them I've created a separate class with the __call__ method. I've done this to have this parameters set in the LambdaLR without further refering to cfg during the training.
In order to avoid branch of if's I've created a dict, which chooses one of the options, and in the function create_scheduler the instance of one of these classes is created, passed to the initializer of LambdaLR scheduler, and this scheduler is then returned.
class LR_Steps:
    def __init__(self, cfg):
        self.lr_base = cfg[&quot;OPTIM&quot;][&quot;BASE_LR&quot;] 
        self.lr_mult = cfg[&quot;OPTIM&quot;][&quot;LR_MULT&quot;] 
        self.steps   =  cfg[&quot;OPTIM&quot;][&quot;STEPS&quot;]
        # the current index in the steps
        self.ind = 0
    
    def __call__(self, epoch):
        if epoch &gt; self.steps[self.ind]:
            self.ind += 1
    
        return self.lr_base * (self.lr_mult ** self.ind)

class LR_Exp:
    def __init__(self, cfg):
        self.lr_base = cfg[&quot;OPTIM&quot;][&quot;BASE_LR&quot;] 
        self.lr_mult = cfg[&quot;OPTIM&quot;][&quot;LR_MULT&quot;] 
    
    def __call__(self, epoch):
        return self.lr_base * (self.lr_mult ** epoch)

class LR_Cos:
    def __init__(self, cfg):
        self.lr_base = cfg[&quot;OPTIM&quot;][&quot;BASE_LR&quot;] 
        self.max_epoch = cfg[&quot;OPTIM&quot;][&quot;MAX_EPOCH&quot;]
    
    def __call__(self, epoch):
        return 0.5 * self.lr_base * (1.0 + np.cos(np.pi * epoch / self.max_epoch))

lr_policy_dict = {
    &quot;cos&quot; : LR_Cos,
    &quot;exp&quot; : LR_Exp,
    &quot;steps&quot; : LR_Steps
}

def create_scheduler(optimizer, cfg):
    return optim.lr_scheduler.LambdaLR(optimizer, lr_policy_dict[cfg[&quot;OPTIM&quot;][&quot;LR_POLICY&quot;]](cfg))

My question is - if there a shorter or safer way to do this? I would also appreciate criticism, concerning the structure of code and functionality.
Overall I rather like this. It's a nice way to think about &quot;configurable functions&quot;. Minor stuff:

Due to order of operations, the parens around self.lr_base * (self.lr_mult ** epoch) are not needed
It's worth splitting out a temporary variable for cfg['OPTIM']
Consider annotating your member variables and method signatures with PEP484 type hints
Decimals in 1.0 are not necessary due to float promotion
Spaces before colons - &quot;cos&quot; : - are not PEP8-standard
You've not shown unit tests. If you don't have them, write them.

If you're interested, a different (not necessarily better) way of thinking about these functions is that they can start as normal, non-class functions, where the first parameter(s) are the configuration needed; then call partial() to get a callable reference that you can then use as if it has only epoch as a parameter from the perspective of  LambdaLR.