File size: 598 Bytes
8e1e3ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | def calc_mean_std(feat, eps=1e-5):
N, C, H, W = feat.shape
feat_reshaped = feat.view(N, C, H * W)
mean = feat_reshaped.mean(dim=2)
var = feat_reshaped.var(dim=2)
std = (var + eps).sqrt()
mean = mean.view(N, C, 1, 1)
std = std.view(N, C, 1, 1)
return mean, std
def adain(content_feat, style_feat):
style_mean, style_std = calc_mean_std(style_feat)
content_mean, content_std = calc_mean_std(content_feat)
normalized_content = (content_feat - content_mean) / content_std
output = normalized_content * style_std + style_mean
return output |