57 lines
2.4 KiB
Python
57 lines
2.4 KiB
Python
import os
|
|
import cv2
|
|
import torch
|
|
import numpy as np
|
|
import torch.nn as nn
|
|
from torch.nn import init
|
|
import torch.nn.functional as F
|
|
from torch.optim import lr_scheduler
|
|
|
|
####################################################################
|
|
# ------------------------- Basic Functions -------------------------
|
|
####################################################################
|
|
|
|
def get_scheduler(optimizer, opts, cur_ep=-1):
|
|
if opts.lr_policy == 'lambda':
|
|
def lambda_rule(ep):
|
|
lr_l = 1.0 - max(0, ep - opts.n_ep_decay) / float(opts.n_ep - opts.n_ep_decay + 1)
|
|
return lr_l
|
|
|
|
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule, last_epoch=cur_ep)
|
|
elif opts.lr_policy == 'step':
|
|
scheduler = lr_scheduler.StepLR(optimizer, step_size=opts.n_ep_decay, gamma=0.1, last_epoch=cur_ep)
|
|
elif opts.lr_policy == 'plateau':
|
|
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
|
|
else:
|
|
return NotImplementedError('learning rate policy [%s] is not implemented', opts.lr_policy)
|
|
return scheduler
|
|
|
|
def init_weights(net, init_type, gain):
|
|
def init_func(m):
|
|
classname = m.__class__.__name__
|
|
if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
|
|
if init_type == 'normal':
|
|
init.normal_(m.weight.data, 0.0, gain)
|
|
elif init_type == 'xavier':
|
|
init.xavier_normal_(m.weight.data, gain=gain)
|
|
elif init_type == 'kaiming':
|
|
init.kaiming_normal_(m.weight.data, a=0, mode='fainplanes')
|
|
elif init_type == 'orthogonal':
|
|
init.orthogonal_(m.weight.data, gain=gain)
|
|
else:
|
|
raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
|
|
if hasattr(m, 'bias') and m.bias is not None:
|
|
init.constant_(m.bias.data, 0.0)
|
|
elif classname.find('BatchNorm2d') != -1:
|
|
init.normal_(m.weight.data, 1.0, gain)
|
|
init.constant_(m.bias.data, 0.0)
|
|
|
|
print('initialize network with %s' % init_type)
|
|
net.apply(init_func)
|
|
|
|
def init_net(net, gpu, init_type='normal', gain=0.02):
|
|
# assert (torch.cuda.is_available())
|
|
net.to(gpu)
|
|
init_weights(net, init_type, gain)
|
|
return net
|