150 lines
5.3 KiB
Python
150 lines
5.3 KiB
Python
from torch import nn
|
|
import torch.nn.functional as F
|
|
|
|
from functools import reduce
|
|
|
|
|
|
class BaseNetwork(nn.Module):
|
|
def __init__(self):
|
|
super(BaseNetwork, self).__init__()
|
|
|
|
def print_network(self):
|
|
if isinstance(self, list):
|
|
self = self[0]
|
|
num_params = 0
|
|
for param in self.parameters():
|
|
num_params += param.numel()
|
|
print(
|
|
"Network [%s] was created. Total number of parameters: %.1f million. "
|
|
"" % (type(self).__name__, num_params / 1000000)
|
|
)
|
|
|
|
def init_weights(self, init_type="normal", gain=0.02):
|
|
"""Initialize network's weights
|
|
init_type: normal | xavier | kaiming | orthogonal
|
|
https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/9451e70673400885567d08a9e97ade2524c700d0/models/networks.py#L39
|
|
"""
|
|
|
|
def init_func(m):
|
|
classname = m.__class__.__name__
|
|
if classname.find("InstanceNorm2d") != -1:
|
|
if hasattr(m, "weight") and m.weight is not None:
|
|
nn.init.constant_(m.weight.data, 1.0)
|
|
if hasattr(m, "bias") and m.bias is not None:
|
|
nn.init.constant_(m.bias.data, 0.0)
|
|
elif hasattr(m, "weight") and (
|
|
classname.find("Conv") != -1 or classname.find("Linear") != -1
|
|
):
|
|
if init_type == "normal":
|
|
nn.init.normal_(m.weight.data, 0.0, gain)
|
|
elif init_type == "xavier":
|
|
nn.init.xavier_normal_(m.weight.data, gain=gain)
|
|
elif init_type == "xavier_uniform":
|
|
nn.init.xavier_uniform_(m.weight.data, gain=1.0)
|
|
elif init_type == "kaiming":
|
|
nn.init.kaiming_normal_(m.weight.data, a=0, mode="fan_in")
|
|
elif init_type == "orthogonal":
|
|
nn.init.orthogonal_(m.weight.data, gain=gain)
|
|
elif init_type == "none": # uses pytorch's default init method
|
|
m.reset_parameters()
|
|
else:
|
|
raise NotImplementedError(
|
|
"initialization method [%s] is not implemented" % init_type
|
|
)
|
|
if hasattr(m, "bias") and m.bias is not None:
|
|
nn.init.constant_(m.bias.data, 0.0)
|
|
|
|
self.apply(init_func)
|
|
|
|
# propagate to children
|
|
for m in self.children():
|
|
if hasattr(m, "init_weights"):
|
|
m.init_weights(init_type, gain)
|
|
|
|
|
|
class Vec2Feat(nn.Module):
|
|
def __init__(self, channel, hidden, kernel_size, stride, padding):
|
|
super(Vec2Feat, self).__init__()
|
|
self.relu = nn.LeakyReLU(0.2, inplace=True)
|
|
c_out = reduce((lambda x, y: x * y), kernel_size) * channel
|
|
self.embedding = nn.Linear(hidden, c_out)
|
|
self.kernel_size = kernel_size
|
|
self.stride = stride
|
|
self.padding = padding
|
|
self.bias_conv = nn.Conv2d(channel, channel, kernel_size=3, stride=1, padding=1)
|
|
|
|
def forward(self, x, t, output_size):
|
|
b_, _, _, _, c_ = x.shape
|
|
x = x.view(b_, -1, c_)
|
|
feat = self.embedding(x)
|
|
b, _, c = feat.size()
|
|
feat = feat.view(b * t, -1, c).permute(0, 2, 1)
|
|
feat = F.fold(
|
|
feat,
|
|
output_size=output_size,
|
|
kernel_size=self.kernel_size,
|
|
stride=self.stride,
|
|
padding=self.padding,
|
|
)
|
|
feat = self.bias_conv(feat)
|
|
return feat
|
|
|
|
|
|
class FusionFeedForward(nn.Module):
|
|
def __init__(self, dim, hidden_dim=1960, t2t_params=None):
|
|
super(FusionFeedForward, self).__init__()
|
|
# We set hidden_dim as a default to 1960
|
|
self.fc1 = nn.Sequential(nn.Linear(dim, hidden_dim))
|
|
self.fc2 = nn.Sequential(nn.GELU(), nn.Linear(hidden_dim, dim))
|
|
assert t2t_params is not None
|
|
self.t2t_params = t2t_params
|
|
self.kernel_shape = reduce(
|
|
(lambda x, y: x * y), t2t_params["kernel_size"]
|
|
) # 49
|
|
|
|
def forward(self, x, output_size):
|
|
n_vecs = 1
|
|
for i, d in enumerate(self.t2t_params["kernel_size"]):
|
|
n_vecs *= int(
|
|
(output_size[i] + 2 * self.t2t_params["padding"][i] - (d - 1) - 1)
|
|
/ self.t2t_params["stride"][i]
|
|
+ 1
|
|
)
|
|
|
|
x = self.fc1(x)
|
|
b, n, c = x.size()
|
|
normalizer = (
|
|
x.new_ones(b, n, self.kernel_shape)
|
|
.view(-1, n_vecs, self.kernel_shape)
|
|
.permute(0, 2, 1)
|
|
)
|
|
normalizer = F.fold(
|
|
normalizer,
|
|
output_size=output_size,
|
|
kernel_size=self.t2t_params["kernel_size"],
|
|
padding=self.t2t_params["padding"],
|
|
stride=self.t2t_params["stride"],
|
|
)
|
|
|
|
x = F.fold(
|
|
x.view(-1, n_vecs, c).permute(0, 2, 1),
|
|
output_size=output_size,
|
|
kernel_size=self.t2t_params["kernel_size"],
|
|
padding=self.t2t_params["padding"],
|
|
stride=self.t2t_params["stride"],
|
|
)
|
|
|
|
x = (
|
|
F.unfold(
|
|
x / normalizer,
|
|
kernel_size=self.t2t_params["kernel_size"],
|
|
padding=self.t2t_params["padding"],
|
|
stride=self.t2t_params["stride"],
|
|
)
|
|
.permute(0, 2, 1)
|
|
.contiguous()
|
|
.view(b, n, c)
|
|
)
|
|
x = self.fc2(x)
|
|
return x
|