373 lines
11 KiB
Python
373 lines
11 KiB
Python
import math
|
|
import numbers
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from einops import rearrange
|
|
from torch.nn import functional as F
|
|
|
|
|
|
# ---------------------------------------------------------------------------------------------------------------------
|
|
# Layer Norm
|
|
def to_3d(x):
|
|
return rearrange(x, "b c h w -> b (h w) c")
|
|
|
|
|
|
def to_4d(x, h, w):
|
|
return rearrange(x, "b (h w) c -> b c h w", h=h, w=w)
|
|
|
|
|
|
class BiasFree_LayerNorm(nn.Module):
|
|
def __init__(self, normalized_shape):
|
|
super(BiasFree_LayerNorm, self).__init__()
|
|
if isinstance(normalized_shape, numbers.Integral):
|
|
normalized_shape = (normalized_shape,)
|
|
normalized_shape = torch.Size(normalized_shape)
|
|
|
|
assert len(normalized_shape) == 1
|
|
|
|
self.weight = nn.Parameter(torch.ones(normalized_shape))
|
|
self.normalized_shape = normalized_shape
|
|
|
|
def forward(self, x):
|
|
sigma = x.var(-1, keepdim=True, unbiased=False)
|
|
return x / torch.sqrt(sigma + 1e-5) * self.weight
|
|
|
|
|
|
class WithBias_LayerNorm(nn.Module):
|
|
def __init__(self, normalized_shape):
|
|
super(WithBias_LayerNorm, self).__init__()
|
|
if isinstance(normalized_shape, numbers.Integral):
|
|
normalized_shape = (normalized_shape,)
|
|
normalized_shape = torch.Size(normalized_shape)
|
|
|
|
assert len(normalized_shape) == 1
|
|
|
|
self.weight = nn.Parameter(torch.ones(normalized_shape))
|
|
self.bias = nn.Parameter(torch.zeros(normalized_shape))
|
|
self.normalized_shape = normalized_shape
|
|
|
|
def forward(self, x):
|
|
mu = x.mean(-1, keepdim=True)
|
|
sigma = x.var(-1, keepdim=True, unbiased=False)
|
|
return (x - mu) / torch.sqrt(sigma + 1e-5) * self.weight + self.bias
|
|
|
|
|
|
class LayerNorm(nn.Module):
|
|
def __init__(self, dim, LayerNorm_type):
|
|
super(LayerNorm, self).__init__()
|
|
if LayerNorm_type == "BiasFree":
|
|
self.body = BiasFree_LayerNorm(dim)
|
|
else:
|
|
self.body = WithBias_LayerNorm(dim)
|
|
|
|
def forward(self, x):
|
|
h, w = x.shape[-2:]
|
|
return to_4d(self.body(to_3d(x)), h, w)
|
|
|
|
|
|
# ---------------------------------------------------------------------------------------------------------------------
|
|
|
|
|
|
# FFN
|
|
class FeedForward(nn.Module):
|
|
def __init__(self, dim, ffn_expansion_factor, bias, input_resolution=None):
|
|
super(FeedForward, self).__init__()
|
|
|
|
self.input_resolution = input_resolution
|
|
self.dim = dim
|
|
self.ffn_expansion_factor = ffn_expansion_factor
|
|
|
|
hidden_features = int(dim * ffn_expansion_factor)
|
|
self.project_in = nn.Conv2d(dim, hidden_features * 2, kernel_size=1, bias=bias)
|
|
self.dwconv = nn.Conv2d(
|
|
hidden_features * 2,
|
|
hidden_features * 2,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1,
|
|
groups=hidden_features * 2,
|
|
bias=bias,
|
|
)
|
|
self.project_out = nn.Conv2d(hidden_features, dim, kernel_size=1, bias=bias)
|
|
|
|
def forward(self, x):
|
|
x = self.project_in(x)
|
|
x1, x2 = self.dwconv(x).chunk(2, dim=1)
|
|
x = F.gelu(x1) * x2
|
|
x = self.project_out(x)
|
|
return x
|
|
|
|
|
|
class NormLayer(nn.Module):
|
|
"""Normalization Layers.
|
|
------------
|
|
# Arguments
|
|
- channels: input channels, for batch norm and instance norm.
|
|
- input_size: input shape without batch size, for layer norm.
|
|
"""
|
|
|
|
def __init__(self, channels, normalize_shape=None, norm_type="bn"):
|
|
super(NormLayer, self).__init__()
|
|
norm_type = norm_type.lower()
|
|
if norm_type == "bn":
|
|
self.norm = nn.BatchNorm2d(channels)
|
|
elif norm_type == "in":
|
|
self.norm = nn.InstanceNorm2d(channels, affine=True)
|
|
elif norm_type == "gn":
|
|
self.norm = nn.GroupNorm(32, channels, affine=True)
|
|
elif norm_type == "pixel":
|
|
self.norm = lambda x: F.normalize(x, p=2, dim=1)
|
|
elif norm_type == "layer":
|
|
self.norm = nn.LayerNorm(normalize_shape)
|
|
elif norm_type == "none":
|
|
self.norm = lambda x: x
|
|
else:
|
|
assert 1 == 0, "Norm type {} not support.".format(norm_type)
|
|
|
|
def forward(self, x):
|
|
return self.norm(x)
|
|
|
|
|
|
class ReluLayer(nn.Module):
|
|
"""Relu Layer.
|
|
------------
|
|
# Arguments
|
|
- relu type: type of relu layer, candidates are
|
|
- ReLU
|
|
- LeakyReLU: default relu slope 0.2
|
|
- PRelu
|
|
- SELU
|
|
- none: direct pass
|
|
"""
|
|
|
|
def __init__(self, channels, relu_type="relu"):
|
|
super(ReluLayer, self).__init__()
|
|
relu_type = relu_type.lower()
|
|
if relu_type == "relu":
|
|
self.func = nn.ReLU(True)
|
|
elif relu_type == "leakyrelu":
|
|
self.func = nn.LeakyReLU(0.2, inplace=True)
|
|
elif relu_type == "prelu":
|
|
self.func = nn.PReLU(channels)
|
|
elif relu_type == "selu":
|
|
self.func = nn.SELU(True)
|
|
elif relu_type == "none":
|
|
self.func = lambda x: x
|
|
else:
|
|
assert 1 == 0, "Relu type {} not support.".format(relu_type)
|
|
|
|
def forward(self, x):
|
|
return self.func(x)
|
|
|
|
|
|
class ConvLayer(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels,
|
|
out_channels,
|
|
kernel_size=3,
|
|
scale="none",
|
|
norm_type="none",
|
|
relu_type="none",
|
|
use_pad=True,
|
|
):
|
|
super(ConvLayer, self).__init__()
|
|
self.use_pad = use_pad
|
|
|
|
bias = True if norm_type in ["pixel", "none"] else False
|
|
stride = 2 if scale == "down" else 1
|
|
|
|
self.scale_func = lambda x: x
|
|
if scale == "up":
|
|
self.scale_func = lambda x: nn.functional.interpolate(
|
|
x, scale_factor=2, mode="nearest"
|
|
)
|
|
|
|
self.reflection_pad = nn.ReflectionPad2d(kernel_size // 2)
|
|
self.conv2d = nn.Conv2d(
|
|
in_channels, out_channels, kernel_size, stride, bias=bias
|
|
)
|
|
|
|
self.relu = ReluLayer(out_channels, relu_type)
|
|
self.norm = NormLayer(out_channels, norm_type=norm_type)
|
|
|
|
def forward(self, x):
|
|
out = self.scale_func(x)
|
|
if self.use_pad:
|
|
out = self.reflection_pad(out)
|
|
out = self.conv2d(out)
|
|
out = self.norm(out)
|
|
out = self.relu(out)
|
|
return out
|
|
|
|
|
|
class ResidualBlock(nn.Module):
|
|
"""
|
|
Residual block recommended in: http://torch.ch/blog/2016/02/04/resnets.html
|
|
------------------
|
|
# Args
|
|
- hg_depth: depth of HourGlassBlock. 0: don't use attention map.
|
|
- use_pmask: whether use previous mask as HourGlassBlock input.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
c_in,
|
|
c_out,
|
|
relu_type="prelu",
|
|
norm_type="bn",
|
|
scale="none",
|
|
hg_depth=2,
|
|
att_name="spar",
|
|
):
|
|
super(ResidualBlock, self).__init__()
|
|
self.c_in = c_in
|
|
self.c_out = c_out
|
|
self.norm_type = norm_type
|
|
self.relu_type = relu_type
|
|
self.hg_depth = hg_depth
|
|
|
|
kwargs = {"norm_type": norm_type, "relu_type": relu_type}
|
|
|
|
if scale == "none" and c_in == c_out:
|
|
self.shortcut_func = lambda x: x
|
|
else:
|
|
self.shortcut_func = ConvLayer(c_in, c_out, 3, scale)
|
|
|
|
self.preact_func = nn.Sequential(
|
|
NormLayer(c_in, norm_type=self.norm_type),
|
|
ReluLayer(c_in, self.relu_type),
|
|
)
|
|
|
|
if scale == "down":
|
|
scales = ["none", "down"]
|
|
elif scale == "up":
|
|
scales = ["up", "none"]
|
|
elif scale == "none":
|
|
scales = ["none", "none"]
|
|
|
|
self.conv1 = ConvLayer(c_in, c_out, 3, scales[0], **kwargs)
|
|
self.conv2 = ConvLayer(
|
|
c_out, c_out, 3, scales[1], norm_type=norm_type, relu_type="none"
|
|
)
|
|
|
|
if att_name.lower() == "spar":
|
|
c_attn = 1
|
|
elif att_name.lower() == "spar3d":
|
|
c_attn = c_out
|
|
else:
|
|
raise Exception("Attention type {} not implemented".format(att_name))
|
|
|
|
self.att_func = HourGlassBlock(self.hg_depth, c_out, c_attn, **kwargs)
|
|
|
|
def forward(self, x):
|
|
identity = self.shortcut_func(x)
|
|
out = self.preact_func(x)
|
|
out = self.conv1(out)
|
|
out = self.conv2(out)
|
|
out = identity + self.att_func(out)
|
|
return out
|
|
|
|
|
|
class HourGlassBlock(nn.Module):
|
|
"""Simplified HourGlass block.
|
|
Reference: https://github.com/1adrianb/face-alignment
|
|
--------------------------
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
depth,
|
|
c_in,
|
|
c_out,
|
|
c_mid=64,
|
|
norm_type="bn",
|
|
relu_type="prelu",
|
|
):
|
|
super(HourGlassBlock, self).__init__()
|
|
self.depth = depth
|
|
self.c_in = c_in
|
|
self.c_mid = c_mid
|
|
self.c_out = c_out
|
|
self.kwargs = {"norm_type": norm_type, "relu_type": relu_type}
|
|
|
|
if self.depth:
|
|
self._generate_network(self.depth)
|
|
self.out_block = nn.Sequential(
|
|
ConvLayer(self.c_mid, self.c_out, norm_type="none", relu_type="none"),
|
|
nn.Sigmoid(),
|
|
)
|
|
|
|
def _generate_network(self, level):
|
|
if level == self.depth:
|
|
c1, c2 = self.c_in, self.c_mid
|
|
else:
|
|
c1, c2 = self.c_mid, self.c_mid
|
|
|
|
self.add_module("b1_" + str(level), ConvLayer(c1, c2, **self.kwargs))
|
|
self.add_module(
|
|
"b2_" + str(level), ConvLayer(c1, c2, scale="down", **self.kwargs)
|
|
)
|
|
if level > 1:
|
|
self._generate_network(level - 1)
|
|
else:
|
|
self.add_module(
|
|
"b2_plus_" + str(level),
|
|
ConvLayer(self.c_mid, self.c_mid, **self.kwargs),
|
|
)
|
|
|
|
self.add_module(
|
|
"b3_" + str(level),
|
|
ConvLayer(self.c_mid, self.c_mid, scale="up", **self.kwargs),
|
|
)
|
|
|
|
def _forward(self, level, in_x):
|
|
up1 = self._modules["b1_" + str(level)](in_x)
|
|
low1 = self._modules["b2_" + str(level)](in_x)
|
|
if level > 1:
|
|
low2 = self._forward(level - 1, low1)
|
|
else:
|
|
low2 = self._modules["b2_plus_" + str(level)](low1)
|
|
|
|
up2 = self._modules["b3_" + str(level)](low2)
|
|
if up1.shape[2:] != up2.shape[2:]:
|
|
up2 = nn.functional.interpolate(up2, up1.shape[2:])
|
|
|
|
return up1 + up2
|
|
|
|
def forward(self, x, pmask=None):
|
|
if self.depth == 0:
|
|
return x
|
|
input_x = x
|
|
x = self._forward(self.depth, x)
|
|
self.att_map = self.out_block(x)
|
|
x = input_x * self.att_map
|
|
return x
|
|
|
|
|
|
class UpSampler(nn.Sequential):
|
|
def __init__(self, scale, n_feats, out_feats):
|
|
m = []
|
|
if scale == 8:
|
|
kernel_size = 3
|
|
elif scale == 16:
|
|
kernel_size = 5
|
|
else:
|
|
kernel_size = 1
|
|
|
|
if (scale & (scale - 1)) == 0: # Is scale = 2^n?
|
|
for _ in range(int(math.log(scale, 2))):
|
|
m.append(
|
|
nn.Conv2d(
|
|
in_channels=n_feats,
|
|
out_channels=4 * out_feats,
|
|
kernel_size=kernel_size,
|
|
stride=1,
|
|
padding=kernel_size // 2,
|
|
)
|
|
)
|
|
m.append(nn.PixelShuffle(upscale_factor=2))
|
|
m.append(nn.PReLU())
|
|
super(UpSampler, self).__init__(*m)
|