File size: 4,375 Bytes
3ce19a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | from math import sqrt
import torch
from torch import nn
import numpy as np
class PixelNorm(nn.Module):
def __init__(self):
super().__init__()
def forward(self, input):
return input / torch.sqrt(torch.mean(input ** 2, dim=1, keepdim=True) + 1e-6)
class FullyConnectedLayer(nn.Module):
def __init__(self, in_features, out_features, bias=True,
activation='linear', lr_multiplier=1, bias_init=0):
super().__init__()
self.activation = activation
self.weight = nn.Parameter(torch.randn([out_features, in_features]) / lr_multiplier)
self.bias = nn.Parameter(torch.full([out_features], np.float32(bias_init))) if bias else None
self.weight_gain = lr_multiplier / np.sqrt(in_features)
self.bias_gain = lr_multiplier
def forward(self, x):
w = self.weight.to(x.dtype) * self.weight_gain
b = self.bias
if b is not None:
b = b.to(x.dtype)
if self.bias_gain != 1:
b = b * self.bias_gain
x = torch.addmm(b.unsqueeze(0), x, w.t())
return x
class EqualLinear(nn.Module):
def __init__(self, in_dim, out_dim):
super().__init__()
linear = nn.Linear(in_dim, out_dim)
linear.bias.data.zero_()
self.linear = linear
def forward(self, input):
return self.linear(input)
def normalize_2nd_moment(x, dim=1, eps=1e-8):
return x * (x.square().mean(dim=dim, keepdim=True) + eps).rsqrt()
class MappingNetowrk(nn.Module):
def __init__(self, code_dim=512, n_mlp=8, mapping_lr_multiplier=1.0):
super().__init__()
layers = [PixelNorm()]
for i in range(n_mlp):
layers.append(FullyConnectedLayer(code_dim, code_dim,
lr_multiplier=mapping_lr_multiplier))
layers.append(nn.LeakyReLU(0.2))
self.style = nn.Sequential(*layers)
def forward(
self,
input,
noise=None,
step=0,
alpha=-1,
mean_style=None,
style_weight=0,
mixing_range=(-1, -1),
):
styles = []
# input = normalize_2nd_moment(input)
if type(input) not in (list, tuple):
input = [input]
for i in input:
x = self.style(i)
styles.append(x)
# batch = input[0].shape[0]
#
# if noise is None:
# noise = []
#
# for i in range(step + 1):
# size = 4 * 2 ** i
# noise.append(torch.randn(batch, 1, size, size, device=input[0].device))
# if mean_style is not None:
# styles_norm = []
#
# for style in styles:
# styles_norm.append(mean_style + style_weight * (style - mean_style))
#
# styles = styles_norm
return styles
def forward_w_trajectory(self, z):
"""One style tensor per MLP block (after each FC + activation), for video/strip viz."""
if type(z) not in (list, tuple):
z = [z]
z0 = z[0]
x = self.style[0](z0)
out_list = []
i = 1
while i < len(self.style):
x = self.style[i](x)
i += 1
if i < len(self.style):
x = self.style[i](x)
i += 1
out_list.append(x)
return out_list
# def mean_style(self, input):
# style = self.style(input).mean(0, keepdim=True)
#
# return style
class AdaptiveInstanceNorm(nn.Module):
def __init__(self, in_channel, style_dim):
super().__init__()
self.norm = nn.InstanceNorm2d(in_channel, eps=1e-3)
self.style = EqualLinear(style_dim, in_channel * 2)
nn.init.zeros_(self.style.linear.bias)
def forward(self, input, style):
style = self.style(style).unsqueeze(2).unsqueeze(3)
gamma, beta = style.chunk(2, 1)
if input.shape[-1] > 1:
out = self.norm(input)
else:
out = input
out = (1 + gamma) * out + beta
return out
class NoiseInjection(nn.Module):
def __init__(self, channel):
super().__init__()
self.weight = nn.Parameter(torch.randn(1, channel, 1, 1), requires_grad=False)
def forward(self, image, spatial_noise):
return image
|