37 lines
927 B
Python
37 lines
927 B
Python
import torch
|
|
import torch.nn as nn
|
|
from model.model_selector import model_selector
|
|
|
|
|
|
def init_model(args, device):
|
|
model = model_selector(args).to(device)
|
|
# Initialize model parameters
|
|
for p in model.parameters():
|
|
if p.dim() > 1:
|
|
nn.init.xavier_uniform_(p)
|
|
else:
|
|
nn.init.uniform_(p)
|
|
return model
|
|
|
|
def init_optimizer(model, args):
|
|
|
|
optimizer = torch.optim.Adam(
|
|
params=model.parameters(),
|
|
lr=args['lr_init'],
|
|
eps=1.0e-8,
|
|
weight_decay=args['weight_decay'],
|
|
amsgrad=False
|
|
)
|
|
|
|
lr_scheduler = None
|
|
if args['lr_decay']:
|
|
lr_decay_steps = [int(step) for step in args['lr_decay_step'].split(',')]
|
|
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
|
|
optimizer=optimizer,
|
|
milestones=lr_decay_steps,
|
|
gamma=args['lr_decay_rate']
|
|
)
|
|
|
|
return optimizer, lr_scheduler
|
|
|