343 lines
14 KiB
Python
343 lines
14 KiB
Python
import os
|
|
# import glob
|
|
# from path import Path
|
|
import math
|
|
import torch
|
|
from torch import nn
|
|
import torch.nn.functional as F
|
|
from torch.optim import lr_scheduler
|
|
# from blocks import *
|
|
|
|
def model_save(ckpt_dir, model, optim_E, optim_S, optim_G, epoch, itr=None):
|
|
if not os.path.exists(ckpt_dir):
|
|
os.makedirs(ckpt_dir)
|
|
|
|
torch.save({'netE': model.netE.state_dict(),
|
|
'netS': model.netS.state_dict(),
|
|
'netG': model.netG.state_dict(),
|
|
'optim_E': optim_E.state_dict(),
|
|
'optim_S': optim_S.state_dict(),
|
|
'optim_G': optim_G.state_dict()},
|
|
'%s/model_iter_%d_epoch_%d.pth' % (ckpt_dir, itr+1, epoch+1))
|
|
|
|
def model_load(checkpoint, ckpt_path, model, optim_E, optim_S, optim_G):
|
|
# if not os.path.exists(ckpt_dir):
|
|
# epoch = -1
|
|
# return model, optim_E, optim_S, optim_G, epoch
|
|
|
|
# ckpt_path = Path(ckpt_dir)
|
|
if checkpoint:
|
|
model_ckpt = ckpt_path + '/' + checkpoint
|
|
else:
|
|
ckpt_lst = ckpt_path.glob('model_iter_*')
|
|
ckpt_lst.sort(key=lambda x: int(x.split('iter_')[1].split('_epoch')[0]))
|
|
model_ckpt = ckpt_lst[-1]
|
|
itr = int(model_ckpt.split('iter_')[1].split('_epoch_')[0])
|
|
epoch = int(model_ckpt.split('iter_')[1].split('_epoch_')[1].split('.')[0])
|
|
print(model_ckpt)
|
|
|
|
dict_model = torch.load(model_ckpt)
|
|
|
|
model.netE.load_state_dict(dict_model['netE'])
|
|
model.netS.load_state_dict(dict_model['netS'])
|
|
model.netG.load_state_dict(dict_model['netG'])
|
|
optim_E.load_state_dict(dict_model['optim_E'])
|
|
optim_S.load_state_dict(dict_model['optim_S'])
|
|
optim_G.load_state_dict(dict_model['optim_G'])
|
|
|
|
return model, optim_E, optim_S, optim_G, epoch, itr
|
|
|
|
def test_model_load(checkpoint, model):
|
|
dict_model = torch.load(checkpoint)
|
|
model.netE.load_state_dict(dict_model['netE'])
|
|
model.netS.load_state_dict(dict_model['netS'])
|
|
model.netG.load_state_dict(dict_model['netG'])
|
|
return model
|
|
|
|
def get_scheduler(optimizer, config):
|
|
if config.lr_policy == 'lambda':
|
|
def lambda_rule(epoch):
|
|
lr_l = 1.0 - max(0, epoch + config.n_epoch - config.n_iter) / float(config.n_iter_decay + 1)
|
|
return lr_l
|
|
|
|
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
|
|
elif config.lr_policy == 'step':
|
|
scheduler = lr_scheduler.StepLR(optimizer, step_size=config.lr_decay_iters, gamma=0.1)
|
|
elif config.lr_policy == 'plateau':
|
|
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
|
|
elif config.lr_policy == 'cosine':
|
|
scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=config.n_iter, eta_min=0)
|
|
else:
|
|
return NotImplementedError('learning rate policy [%s] is not implemented', config.lr_policy)
|
|
return scheduler
|
|
|
|
def update_learning_rate(scheduler, optimizer):
|
|
scheduler.step()
|
|
lr = optimizer.param_groups[0]['lr']
|
|
print('learning rate = %.7f' % lr)
|
|
|
|
class Oct_Conv_aftup(nn.Module):
|
|
def __init__(self, in_channels, out_channels, kernel_size, stride, padding, pad_type, alpha_in, alpha_out):
|
|
super(Oct_Conv_aftup, self).__init__()
|
|
lf_in = int(in_channels*alpha_in)
|
|
lf_out = int(out_channels*alpha_out)
|
|
hf_in = in_channels - lf_in
|
|
hf_out = out_channels - lf_out
|
|
|
|
self.conv_h = nn.Conv2d(in_channels=hf_in, out_channels=hf_out, kernel_size=kernel_size, stride=stride, padding=padding, padding_mode=pad_type)
|
|
self.conv_l = nn.Conv2d(in_channels=lf_in, out_channels=lf_out, kernel_size=kernel_size, stride=stride, padding=padding, padding_mode=pad_type)
|
|
|
|
def forward(self, x):
|
|
hf, lf = x
|
|
hf = self.conv_h(hf)
|
|
lf = self.conv_l(lf)
|
|
return hf, lf
|
|
|
|
class Oct_conv_reLU(nn.ReLU):
|
|
def forward(self, x):
|
|
hf, lf = x
|
|
hf = super(Oct_conv_reLU, self).forward(hf)
|
|
lf = super(Oct_conv_reLU, self).forward(lf)
|
|
return hf, lf
|
|
|
|
class Oct_conv_lreLU(nn.LeakyReLU):
|
|
def forward(self, x):
|
|
hf, lf = x
|
|
hf = super(Oct_conv_lreLU, self).forward(hf)
|
|
lf = super(Oct_conv_lreLU, self).forward(lf)
|
|
return hf, lf
|
|
|
|
class Oct_conv_up(nn.Upsample):
|
|
def forward(self, x):
|
|
hf, lf = x
|
|
hf = super(Oct_conv_up, self).forward(hf)
|
|
lf = super(Oct_conv_up, self).forward(lf)
|
|
return hf, lf
|
|
|
|
|
|
############## Encoder ##############
|
|
class OctConv(nn.Module):
|
|
def __init__(self, in_channels, out_channels, kernel_size, stride=1,
|
|
padding=0, groups=1, pad_type='reflect', alpha_in=0.5, alpha_out=0.5, type='normal', freq_ratio = [1, 1]):
|
|
super(OctConv, self).__init__()
|
|
self.kernel_size = kernel_size
|
|
self.stride = stride
|
|
self.type = type
|
|
self.alpha_in = alpha_in
|
|
self.alpha_out = alpha_out
|
|
self.freq_ratio = freq_ratio
|
|
|
|
hf_ch_in = int(in_channels * (1 - self.alpha_in))
|
|
hf_ch_out = int(out_channels * (1 -self. alpha_out))
|
|
lf_ch_in = in_channels - hf_ch_in
|
|
lf_ch_out = out_channels - hf_ch_out
|
|
|
|
self.avg_pool = nn.AvgPool2d(kernel_size=2, stride=2)
|
|
self.upsample = nn.Upsample(scale_factor=2)
|
|
|
|
self.is_dw = groups == in_channels
|
|
|
|
if type == 'first':
|
|
self.convh = nn.Conv2d(in_channels, hf_ch_out, kernel_size=kernel_size,
|
|
stride=stride, padding=padding, padding_mode=pad_type, bias = False)
|
|
self.convl = nn.Conv2d(in_channels, lf_ch_out,
|
|
kernel_size=kernel_size, stride=stride, padding=padding, padding_mode=pad_type, bias=False)
|
|
elif type == 'last':
|
|
self.convh = nn.Conv2d(hf_ch_in, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, padding_mode=pad_type, bias=False)
|
|
self.convl = nn.Conv2d(lf_ch_in, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, padding_mode=pad_type, bias=False)
|
|
else:
|
|
self.L2L = nn.Conv2d(
|
|
lf_ch_in, lf_ch_out,
|
|
kernel_size=kernel_size, stride=stride, padding=padding, groups=math.ceil(alpha_in * groups), padding_mode=pad_type, bias=False
|
|
)
|
|
if self.is_dw:
|
|
self.L2H = None
|
|
self.H2L = None
|
|
else:
|
|
self.L2H = nn.Conv2d(
|
|
lf_ch_in, hf_ch_out,
|
|
kernel_size=kernel_size, stride=stride, padding=padding, groups=groups, padding_mode=pad_type, bias=False
|
|
)
|
|
self.H2L = nn.Conv2d(
|
|
hf_ch_in, lf_ch_out,
|
|
kernel_size=kernel_size, stride=stride, padding=padding, groups=groups, padding_mode=pad_type, bias=False
|
|
)
|
|
self.H2H = nn.Conv2d(
|
|
hf_ch_in, hf_ch_out,
|
|
kernel_size=kernel_size, stride=stride, padding=padding, groups=math.ceil(groups - alpha_in * groups), padding_mode=pad_type, bias=False
|
|
)
|
|
|
|
def forward(self, x):
|
|
if self.type == 'first':
|
|
hf = self.convh(x)
|
|
lf = self.avg_pool(x)
|
|
lf = self.convl(lf)
|
|
return hf, lf
|
|
elif self.type == 'last':
|
|
hf, lf = x
|
|
out_h = self.convh(hf)
|
|
out_l = self.convl(self.upsample(lf))
|
|
output = out_h * self.freq_ratio[0] + out_l * self.freq_ratio[1]
|
|
return output, out_h, out_l
|
|
else:
|
|
hf, lf = x
|
|
if self.is_dw:
|
|
hf, lf = self.H2H(hf), self.L2L(lf)
|
|
else:
|
|
hf, lf = self.H2H(hf) + self.L2H(self.upsample(lf)), self.L2L(lf) + self.H2L(self.avg_pool(hf))
|
|
return hf, lf
|
|
|
|
|
|
############## Decoder ##############
|
|
class AdaOctConv(nn.Module):
|
|
def __init__(self, in_channels, out_channels, group_div, style_channels, kernel_size,
|
|
stride, padding, oct_groups, alpha_in, alpha_out, type='normal'):
|
|
super(AdaOctConv, self).__init__()
|
|
self.in_channels = in_channels
|
|
self.alpha_in = alpha_in
|
|
self.alpha_out = alpha_out
|
|
self.type = type
|
|
|
|
h_in = int(in_channels * (1 - self.alpha_in))
|
|
l_in = in_channels - h_in
|
|
|
|
n_groups_h = h_in // group_div
|
|
n_groups_l = l_in // group_div
|
|
|
|
style_channels_h = int(style_channels * (1 - self.alpha_in))
|
|
style_channels_l = int(style_channels - style_channels_h)
|
|
|
|
kernel_size_h = kernel_size[0]
|
|
kernel_size_l = kernel_size[1]
|
|
kernel_size_A = kernel_size[2]
|
|
|
|
self.kernelPredictor_h = KernelPredictor(in_channels=h_in,
|
|
out_channels=h_in,
|
|
n_groups=n_groups_h,
|
|
style_channels=style_channels_h,
|
|
kernel_size=kernel_size_h)
|
|
self.kernelPredictor_l = KernelPredictor(in_channels=l_in,
|
|
out_channels=l_in,
|
|
n_groups=n_groups_l,
|
|
style_channels=style_channels_l,
|
|
kernel_size=kernel_size_l)
|
|
|
|
self.AdaConv_h = AdaConv2d(in_channels=h_in, out_channels=h_in, n_groups=n_groups_h)
|
|
self.AdaConv_l = AdaConv2d(in_channels=l_in, out_channels=l_in, n_groups=n_groups_l)
|
|
|
|
self.OctConv = OctConv(in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
kernel_size=kernel_size_A, stride=stride, padding=padding, groups=oct_groups,
|
|
alpha_in=alpha_in, alpha_out=alpha_out, type=type)
|
|
|
|
self.relu = Oct_conv_lreLU()
|
|
|
|
def forward(self, content, style, cond='train'):
|
|
c_hf, c_lf = content
|
|
s_hf, s_lf = style
|
|
h_w_spatial, h_w_pointwise, h_bias = self.kernelPredictor_h(s_hf)
|
|
l_w_spatial, l_w_pointwise, l_bias = self.kernelPredictor_l(s_lf)
|
|
|
|
if cond == 'train':
|
|
output_h = self.AdaConv_h(c_hf, h_w_spatial, h_w_pointwise, h_bias)
|
|
output_l = self.AdaConv_l(c_lf, l_w_spatial, l_w_pointwise, l_bias)
|
|
output = output_h, output_l
|
|
|
|
output = self.relu(output)
|
|
|
|
output = self.OctConv(output)
|
|
if self.type != 'last':
|
|
output = self.relu(output)
|
|
return output
|
|
|
|
if cond == 'test':
|
|
output_h = self.AdaConv_h(c_hf, h_w_spatial, h_w_pointwise, h_bias)
|
|
output_l = self.AdaConv_l(c_lf, l_w_spatial, l_w_pointwise, l_bias)
|
|
output = output_h, output_l
|
|
output = self.relu(output)
|
|
output = self.OctConv(output)
|
|
if self.type != 'last':
|
|
output = self.relu(output)
|
|
return output
|
|
|
|
class KernelPredictor(nn.Module):
|
|
def __init__(self, in_channels, out_channels, n_groups, style_channels, kernel_size):
|
|
super(KernelPredictor, self).__init__()
|
|
|
|
self.in_channels = in_channels
|
|
self.out_channels = out_channels
|
|
self.n_groups = n_groups
|
|
self.w_channels = style_channels
|
|
self.kernel_size = kernel_size
|
|
|
|
padding = (kernel_size - 1) / 2
|
|
self.spatial = nn.Conv2d(style_channels,
|
|
in_channels * out_channels // n_groups,
|
|
kernel_size=kernel_size,
|
|
padding=(math.ceil(padding), math.ceil(padding)),
|
|
padding_mode='reflect')
|
|
self.pointwise = nn.Sequential(
|
|
nn.AdaptiveAvgPool2d((1, 1)),
|
|
nn.Conv2d(style_channels,
|
|
out_channels * out_channels // n_groups,
|
|
kernel_size=1)
|
|
)
|
|
self.bias = nn.Sequential(
|
|
nn.AdaptiveAvgPool2d((1, 1)),
|
|
nn.Conv2d(style_channels,
|
|
out_channels,
|
|
kernel_size=1)
|
|
)
|
|
|
|
def forward(self, w):
|
|
w_spatial = self.spatial(w)
|
|
w_spatial = w_spatial.reshape(len(w),
|
|
self.out_channels,
|
|
self.in_channels // self.n_groups,
|
|
self.kernel_size, self.kernel_size)
|
|
|
|
w_pointwise = self.pointwise(w)
|
|
w_pointwise = w_pointwise.reshape(len(w),
|
|
self.out_channels,
|
|
self.out_channels // self.n_groups,
|
|
1, 1)
|
|
bias = self.bias(w)
|
|
bias = bias.reshape(len(w), self.out_channels)
|
|
return w_spatial, w_pointwise, bias
|
|
|
|
class AdaConv2d(nn.Module):
|
|
def __init__(self, in_channels, out_channels, kernel_size=3, n_groups=None):
|
|
super(AdaConv2d, self).__init__()
|
|
self.n_groups = in_channels if n_groups is None else n_groups
|
|
self.in_channels = in_channels
|
|
self.out_channels = out_channels
|
|
|
|
padding = (kernel_size - 1) / 2
|
|
self.conv = nn.Conv2d(in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
kernel_size=(kernel_size, kernel_size),
|
|
padding=(math.ceil(padding), math.floor(padding)),
|
|
padding_mode='reflect')
|
|
|
|
def forward(self, x, w_spatial, w_pointwise, bias):
|
|
assert len(x) == len(w_spatial) == len(w_pointwise) == len(bias)
|
|
x = F.instance_norm(x)
|
|
|
|
ys = []
|
|
for i in range(len(x)):
|
|
y = self.forward_single(x[i:i+1], w_spatial[i], w_pointwise[i], bias[i])
|
|
ys.append(y)
|
|
ys = torch.cat(ys, dim=0)
|
|
|
|
ys = self.conv(ys)
|
|
return ys
|
|
|
|
def forward_single(self, x, w_spatial, w_pointwise, bias):
|
|
assert w_spatial.size(-1) == w_spatial.size(-2)
|
|
padding = (w_spatial.size(-1) - 1) / 2
|
|
pad = (math.ceil(padding), math.floor(padding), math.ceil(padding), math.floor(padding))
|
|
|
|
x = F.pad(x, pad=pad, mode='reflect')
|
|
x = F.conv2d(x, w_spatial, groups=self.n_groups)
|
|
x = F.conv2d(x, w_pointwise, groups=self.n_groups, bias=bias)
|
|
return x |