utils.py

import torch.optim as optim
import torch.nn.init as init
import torch


def parse_optimizer(parser):
    opt_parser = parser.add_argument_group()
    opt_parser.add_argument('--opt', dest='opt', type=str, help='Type of optimizer')
    opt_parser.add_argument('--opt-scheduler',
                            dest='opt_scheduler',
                            type=str,
                            help='Type of optimizer scheduler. By default none')
    opt_parser.add_argument('--opt-restart',
                            dest='opt_restart',
                            type=int,
                            help='Number of epochs before restart (by default set to 0 which means no restart)')
    opt_parser.add_argument('--opt-decay-step', dest='opt_decay_step', type=int, help='Number of epochs before decay')
    opt_parser.add_argument('--opt-decay-rate', dest='opt_decay_rate', type=float, help='Learning rate decay ratio')
    opt_parser.add_argument('--lr', dest='lr', type=float, help='Learning rate.')
    opt_parser.add_argument('--clip', dest='clip', type=float, help='Gradient clipping.')
    opt_parser.add_argument('--weight_decay', type=float, help='Optimizer weight decay.')


def build_optimizer(args, params):
    weight_decay = args.weight_decay
    filter_fn = filter(lambda p: p.requires_grad, params)
    if args.opt == 'adam':
        optimizer = optim.Adam(filter_fn, lr=args.lr, weight_decay=weight_decay)
    elif args.opt == 'sgd':
        optimizer = optim.SGD(filter_fn, lr=args.lr, momentum=0.95, weight_decay=weight_decay)
    elif args.opt == 'rmsprop':
        optimizer = optim.RMSprop(filter_fn, lr=args.lr, weight_decay=weight_decay)
    elif args.opt == 'adagrad':
        optimizer = optim.Adagrad(filter_fn, lr=args.lr, weight_decay=weight_decay)
    if args.opt_scheduler == 'none':
        return None, optimizer
    elif args.opt_scheduler == 'step':
        scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=args.opt_decay_step, gamma=args.opt_decay_rate)
    elif args.opt_scheduler == 'cos':
        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.opt_restart)
    return scheduler, optimizer


def init_weight_(weight):
    """
    Initialize a weighting tensor
    """
    init.xavier_uniform_(weight)