55 lines
2.0 KiB
Python
55 lines
2.0 KiB
Python
from torch.optim import lr_scheduler
|
|
|
|
import torch.nn as nn
|
|
import torch
|
|
|
|
######################################################### training utils##########################################################
|
|
|
|
def get_scheduler(optimizer, niter,niter_decay,lr_policy='lambda',lr_decay_iters=50):
|
|
'''
|
|
scheduler in training stage
|
|
'''
|
|
if lr_policy == 'lambda':
|
|
def lambda_rule(epoch):
|
|
lr_l = 1.0 - max(0, epoch - niter) / float(niter_decay + 1)
|
|
return lr_l
|
|
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
|
|
elif lr_policy == 'step':
|
|
scheduler = lr_scheduler.StepLR(optimizer, step_size=lr_decay_iters, gamma=0.1)
|
|
elif lr_policy == 'plateau':
|
|
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
|
|
elif lr_policy == 'cosine':
|
|
scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=niter, eta_min=0)
|
|
else:
|
|
return NotImplementedError('learning rate policy [%s] is not implemented', lr_policy)
|
|
return scheduler
|
|
|
|
def update_learning_rate(scheduler, optimizer):
|
|
scheduler.step()
|
|
lr = optimizer.param_groups[0]['lr']
|
|
print('learning rate = %.7f' % lr)
|
|
|
|
class GANLoss(nn.Module):
|
|
'''
|
|
GAN loss
|
|
'''
|
|
def __init__(self, use_lsgan=True, target_real_label=1.0, target_fake_label=0.0):
|
|
super(GANLoss, self).__init__()
|
|
self.register_buffer('real_label', torch.tensor(target_real_label))
|
|
self.register_buffer('fake_label', torch.tensor(target_fake_label))
|
|
if use_lsgan:
|
|
self.loss = nn.MSELoss()
|
|
else:
|
|
self.loss = nn.BCELoss()
|
|
|
|
def get_target_tensor(self, input, target_is_real):
|
|
if target_is_real:
|
|
target_tensor = self.real_label
|
|
else:
|
|
target_tensor = self.fake_label
|
|
return target_tensor.expand_as(input)
|
|
|
|
def forward(self, input, target_is_real):
|
|
target_tensor = self.get_target_tensor(input, target_is_real)
|
|
return self.loss(input, target_tensor)
|