|
import torch |
|
|
|
def calc_mean_std(feat, eps=1e-5): |
|
|
|
size = feat.size() |
|
assert (len(size) == 4) |
|
N, C = size[:2] |
|
feat_var = feat.view(N, C, -1).var(dim=2) + eps |
|
feat_std = feat_var.sqrt().view(N, C, 1, 1) |
|
feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1) |
|
return feat_mean, feat_std |
|
|
|
def adaptive_instance_normalization(content_feat, style_feat): |
|
assert (content_feat.size()[:2] == style_feat.size()[:2]) |
|
size = content_feat.size() |
|
style_mean, style_std = calc_mean_std(style_feat) |
|
content_mean, content_std = calc_mean_std(content_feat) |
|
normalized_feat = (content_feat - content_mean.expand( |
|
size)) / content_std.expand(size) |
|
|
|
return normalized_feat * style_std.expand(size) + style_mean.expand(size) |
|
|
|
def _calc_feat_flatten_mean_std(feat): |
|
|
|
assert (feat.size()[0] == 3) |
|
assert (isinstance(feat, torch.FloatTensor)) |
|
feat_flatten = feat.view(3, -1) |
|
mean = feat_flatten.mean(dim=-1, keepdim=True) |
|
std = feat_flatten.std(dim=-1, keepdim=True) |
|
return feat_flatten, mean, std |
|
|
|
def _mat_sqrt(x): |
|
U, D, V = torch.svd(x) |
|
return torch.mm(torch.mm(U, D.pow(0.5).diag()), V.t()) |
|
|
|
def coral(source, target): |
|
|
|
|
|
source_f, source_f_mean, source_f_std = _calc_feat_flatten_mean_std(source) |
|
source_f_norm = (source_f - source_f_mean.expand_as( |
|
source_f)) / source_f_std.expand_as(source_f) |
|
source_f_cov_eye = \ |
|
torch.mm(source_f_norm, source_f_norm.t()) + torch.eye(3) |
|
|
|
target_f, target_f_mean, target_f_std = _calc_feat_flatten_mean_std(target) |
|
target_f_norm = (target_f - target_f_mean.expand_as( |
|
target_f)) / target_f_std.expand_as(target_f) |
|
target_f_cov_eye = \ |
|
torch.mm(target_f_norm, target_f_norm.t()) + torch.eye(3) |
|
|
|
source_f_norm_transfer = torch.mm( |
|
_mat_sqrt(target_f_cov_eye), |
|
torch.mm(torch.inverse(_mat_sqrt(source_f_cov_eye)), |
|
source_f_norm) |
|
) |
|
|
|
source_f_transfer = source_f_norm_transfer * \ |
|
target_f_std.expand_as(source_f_norm) + \ |
|
target_f_mean.expand_as(source_f_norm) |
|
|
|
return source_f_transfer.view(source.size()) |