StyleSwin: Transformer-based GAN for High-resolution Image Generation论文和代码
论文:
论文地址：StyleSwin: Transformer-Based GAN for High-Resolution Image Generation (thecvf.com)
代码地址： https://github.com/microsoft/StyleSwin
前言：
StyleSwin模型框架是style-based generator上进行修改的，以SwinTransformer为基本框架，在看这篇论文之前，建议先看styleGAN、styleGAN2、SwinTransformer的论文，便于理解本论文中所提出的模型原理。该论文代码也是基于styleGAN2和SwinTransformer的代码进行修改的。
内容：
论文是用一种纯Transformer的生成对抗网络来生成高分辨率的图像。生成器框架如图（b）所示，框架原理理解如下：
整个generator由多个generator blocks和tRGB组成
左分支：一个潜在变量通过由8个全连接层组成的映射网络映射输出空间向量w（维度和z一样，512*1），将latent code 转换得到w后，向量w经可学习的仿射变换A（实际上是全连接层），分别送入右分支的合成网络，进行控制特征。
Mapping Network作用：为输入向量z的特征解耦提供一条学习通道，由于z是符合均匀分布或者高斯分布的随机变量，所以变量之间的耦合性比较大。需要将latent code z进行解耦，才能更好的进行后续操作，来改变其特征。如果仅通过z来控制视觉特征，那么其能力十分有限，因为它必须遵循训练数据的概率密度。比如数据集中长头发的人很常见，那么更多的输入值便会映射到该特征上，那么z中其他变量也会向着该值靠近、无法更好地映射其他特征。所以通过Mapping network，该模型可以生成一个不必遵循训练数据分布地向量w，减少了特征之间的相关性，完成解耦。
主（右）分支（generator）：首先输入一个常量大小4*4*512，经过正弦位置编码（绝对位置）进入generator block,输入到自适应归一化（AdIN）和残差连接，然后注入样式风格，特征图输入到双重注意力层（Double Attn)，给一个相对位置编码，在Double Attn中进行W-MSA和SW-MSA计算，计算完成后特征图进入下一个自适应归一化（AdIN）和残差连接，进行样式风格注入，经过MLP层输出特征图，特征图经上采样加上绝对位置向量进入下一个generator block，每次const和 w 进入一个generator block后，最后输出的feature map作为输入进入下一个generator block，同时会输入到tRGB层中输出一张RGB图像累加到历史的累加RGB图像上。最后，经过不断上采样形成高分辨率的特征图，通过累加的RGB图片作为生成图片（1024*1024）输出。
利用常量输入代替传统的输入，一个是防止不合理的输入生成错误的图像；另一个是有利于减少特征纠缠。
AdIN：基于AdIN可以自调制生成器实现任意图像风格的转换。
特征图的均值和方差带有图像的风格信息，所以在AdIN层中 ，特征图减去自己的均值除以方差，去掉自己的风格。再乘上新风格的方差加上均值，以实现转换的目的。
Double Attn中M-MSA/SM-MSA的计算方法：给定l层的输入特征图Xl，连续的swin transformer操作如下：
代码:
参考：https://zhuanlan.zhihu.com/p/391980528
论文代码是基于styleGAN2和SwinTransformer的代码进行修改的，所以可以结合styleGAN2和SwinTransforme的代码进行理解。
代码是基于Pytorch框架写的
generator.py
ToRGB： ToRGB层由一个全连接层，一个自定义卷积层和一个激活函数层组成，包含可学习参数bias
class ToRGB(nn.Module):
def __init__(self, in_channel, upsample=True, resolution=None, blur_kernel=[1, 3, 3, 1]):
super().__init__()
self.is_upsample = upsample
self.resolution = resolution
if upsample:
self.upsample = Upsample(blur_kernel)
# 将输入feature map的channel变为3（即rgb三通道图），kernel_size=1x1，不变空间维度
self.conv = nn.Conv2d(in_channel, 3, kernel_size=1)
self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1))
def forward(self, input, skip=None):
out = self.conv(input)
out = out + self.bias
if skip is not None:
if self.is_upsample:
skip = self.upsample(skip)
out = out + skip
return out
def flops(self):
m = self.conv
kernel_ops = torch.zeros(m.weight.size()[2:]).numel() # Kw x Kh
bias_ops = 1
# N x Cout x H x W x (Cin x Kw x Kh + bias)
flops = 1 * self.resolution * self.resolution * 3 * (m.in_channels // m.groups * kernel_ops + bias_ops)
if self.is_upsample:
# there is a conv used in upsample
w_shape = (1, 1, 4, 4)
kernel_ops = torch.zeros(w_shape[2:]).numel() # Kw x Kh
# N x Cout x H x W x (Cin x Kw x Kh + bias)
flops = 1 * 3 * (2 * self.resolution + 3) * (2 *self.resolution + 3) * (3 * kernel_ops)
return flopsMLP：构建MLP模块，该模块在SwinTransformer中的末尾位置，用于对SW-MSA形成的全局注意力进行整合。该模块由输入层、一层隐藏层和输出层组成。输入特征经过全连接层，由激活函数GELU函数进行激活，并添加Drop用于防止过拟合，生成隐藏特征。然后，隐藏特征经过第二个全连接层，再添加Drop层，生成了MLP的输出层特征。
class Mlp(nn.Module):
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.hidden_features = hidden_features
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return xwindow_partition: 窗口分割函数构建，第一个输入为图像Batch,对应的尺寸为（B,H,W,C),其中B为batch_size,H和W代表图像的高和宽，C为通道数，注意，这里的通道数为3，而torch直接读取的图像中通道数为1，因此要对图片做预处理。第二个输入为窗口的大小。
def window_partition(x, window_size):
"""
Args:
x: (B, H, W, C)
window_size (int): window size
Returns:
windows: (num_windows*B, window_size, window_size, C)
"""
B, H, W, C = x.shape
x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
return windows首先将输入的特行形状调整为（B, H // window_size, window_size, W // window_size, window_size, C），6个维度参数，第2个维度和第4个维度参数为window_size，说明此处的窗口分割是正方形的形状。使用permute()来进行调换第2个和第3个特征维度，调整后的形状为（B, H // window_size, W // window_size, window_size, window_size, C）,便于后面变形。然后使用view()操作进行最后的变形，特征形状变为（B*H*W//window_size//window_size,window_size,window_size,C），4个维度。
view()：该函数操作，相当于将原来输入的每个图像(H,W,C)处理成(H*W//windows_size//windows_size,windows_size,,windows_size,C)的窗口batch，其中第0维表示窗口的数量，最后再将batch中的所有窗口在第0维进行连接，得到（B*H*W//windows_size//windows_size,windows_size,windows_size,C）。
window_reverse：
def window_reverse(windows, window_size, H, W):
"""
Args:
windows: (num_windows*B, window_size, window_size, C)
window_size (int): Window size
H (int): Height of image
W (int): Width of image
Returns:
x: (B, H, W, C)
"""
B = int(windows.shape[0] / (H * W / window_size / window_size))
x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
return xWindows_reverse()与上面的函数操作相反，将若干个分割好的小窗口，还原成一个图像batch。对于原来batch_size的计算，可以参考上面背景色部分，其余的步骤就是对windows_partition函数的还原操作了，最后尺寸为（B, H, W, C)，该函数用于对局部注意力Batch还原为与图像相同batch的尺寸。
Windows_Attention：
class WindowAttention(nn.Module):
r""" Window based multi-head self attention (W-MSA) module with relative position bias.
It supports both of shifted and non-shifted window.
Args:
dim (int): Number of input channels.
window_size (tuple[int]): The height and width of the window.
num_heads (int): Number of attention heads.
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
"""参数声明：输入的窗口尺寸为（nW*B, window_size*window_size, C），
im(int)：输入通道数C，
window_size(tuple[int])：窗口尺寸的高和宽，
num_heads(int)：注意力的head数量，每个head分别对QKV求若干次attention，然后对所有的head进行concat形成了最终的multihead attention，
qkv_bias(bool,optional)：计算qkv时是否使用偏置项
qk_scale (float | None, optional): q的放缩因子，若未设置则使用head_dim ** -0.5
attn_drop (float, optional): 注意力的dropout
定义先略过，根据需要在回来看
def __init__(self, dim, window_size, num_heads, qk_scale=None, attn_drop=0.):
super().__init__()
self.dim = dim
self.window_size = window_size # Wh, Ww
self.num_heads = num_heads
head_dim = dim // num_heads
self.head_dim = head_dim
self.scale = qk_scale or head_dim ** -0.5
# define a parameter table of relative position bias
self.relative_position_bias_table = nn.Parameter(
torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH
# get pair-wise relative position index for each token inside the window
coords_h = torch.arange(self.window_size[0])
coords_w = torch.arange(self.window_size[1])
coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww
coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww
relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0
relative_coords[:, :, 1] += self.window_size[1] - 1
relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
self.register_buffer("relative_position_index", relative_position_index)
trunc_normal_(self.relative_position_bias_table, std=.02)
self.attn_drop = nn.Dropout(attn_drop)
self.softmax = nn.Softmax(dim=-1)forward()：
def forward(self, q, k, v, mask=None):
"""
Args:
q: queries with shape of (num_windows*B, N, C)
k: keys with shape of (num_windows*B, N, C)
v: values with shape of (num_windows*B, N, C)
mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
"""
B_, N, C = q.shape
q = q.reshape(B_, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
k = k.reshape(B_, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
v = v.reshape(B_, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
q = q * self.scale
attn = (q @ k.transpose(-2, -1))
relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH
relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww
attn = attn + relative_position_bias.unsqueeze(0)
if mask is not None:
nW = mask.shape[0]
attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
attn = attn.view(-1, self.num_heads, N, N)
attn = self.softmax(attn)
else:
attn = self.softmax(attn)
attn = self.attn_drop(attn)
x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
return x
①获取输入特征的尺寸，对应（nW*B,window_size*window_size,C) 记为（B_,N,C)
②使用全连接层计算qkv，使用全连接层计算qkv，得到（nW*B, window_size*window_size, 3C），进一步通过变形和换位，转换形状（B_N，self.num_size, N, C//self.num_size），便于后续矩阵乘法
③分别提取出q、k、v，并对q的值进行加权处理
④将k进行转置后，与q进行矩阵乘法得到attention
⑤计算relative_position_bias，并和attention相加
⑥当计算SW-MSA时，mask不为None，此时需要将attention转为与mask相同的形状，相加后再还原原形状
⑦对attention进行随机的dropout，然后与v进行矩阵乘法，并转换形状为(B_, N, C)
def extra_repr(self) -> str:
return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}'
def flops(self, N):
# calculate flops for 1 window with token length of N
flops = 0
# qkv = self.qkv(x)
flops += N * self.dim * 3 * self.dim
# attn = (q @ k.transpose(-2, -1))
flops += self.num_heads * N * (self.dim // self.num_heads) * N
# x = (attn @ v)
flops += self.num_heads * N * N * (self.dim // self.num_heads)
# x = self.proj(x)
flops += N * self.dim * self.dim
return flopsAdIN：自适应实例归一化层
class AdaptiveInstanceNorm(nn.Module): # AdIN
def __init__(self, in_channel, style_dim):
super().__init__()
self.norm = nn.InstanceNorm1d(in_channel) # 创建IN层
self.style = EqualLinear(style_dim, in_channel * 2) # 从w向量变成AdIN层系数
def forward(self, input, style):
# print("AdIN style input="+str(style.shape)) # 默认值，风格向量长度512
# 输入style为风格向量，长度为512；经过self.style得到输出风格矩阵，通道数等于输入通道
style = self.style(style).unsqueeze(-1)
gamma, beta = style.chunk(2, 1) # 获得缩放和偏置系数，按1轴（通道）分为2部分
# print("AdIN style input="+str(style.shape))
out = self.norm(input) # IN归一化
out = gamma * out + beta
return out
StyleBasicLayer()：是StyleSwinTransformerBlock的基本组成
class StyleBasicLayer(nn.Module):
""" A basic StyleSwin layer for one stage.
Args:
dim (int): Number of input channels.
input_resolution (tuple[int]): Input resolution.
depth (int): Number of blocks.
num_heads (int): Number of attention heads.
window_size (int): Local window size.
out_dim (int): Number of output channels.
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
drop (float, optional): Dropout rate. Default: 0.0
attn_drop (float, optional): Attention dropout rate. Default: 0.0
upsample (nn.Module | None, optional): Upsample layer at the end of the layer. Default: None
use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
style_dim (int): Dimension of style vector.
"""输入的参数：其中有一大部分是进一步传参到StykeSwinTransformerBlock（）的，会在下面介绍。
depth (int): 模型中block的重复次数
upsample (nn.Module | None, optional)：是否在blocks模块后面使用上采样
use_checkpoint (bool)：是否保存checkpoint
style_dim(int)：样式向量的维度
def __init__(self, dim, input_resolution, depth, num_heads, window_size, out_dim=None,
mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., upsample=None,
use_checkpoint=False, style_dim=512):
super().__init__()
self.dim = dim
self.input_resolution = input_resolution
self.depth = depth
self.use_checkpoint = use_checkpoint
# build blocks
self.blocks = nn.ModuleList([
StyleSwinTransformerBlock(dim=dim, input_resolution=input_resolution,
num_heads=num_heads, window_size=window_size,
mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
drop=drop, attn_drop=attn_drop, style_dim=style_dim)
for _ in range(depth)])
if upsample is not None:
self.upsample = upsample(input_resolution, dim=dim, out_dim=out_dim)
else:
self.upsample = None初始化函数中，分别定义了blocks和upsample，其中，blocks由一些列的循环的StyleSwinTransformerBlock组成，循环的次数由传参depth决定，由此计算“局部→全局注意力”
upsample用于对特征进行上采样
def forward(self, x, latent1, latent2):
if self.use_checkpoint:
x = checkpoint.checkpoint(self.blocks[0], x, latent1)
x = checkpoint.checkpoint(self.blocks[1], x, latent2)
else:
x = self.blocks[0](x, latent1)
x = self.blocks[1](x, latent2)
if self.upsample is not None:
x = self.upsample(x)
return x
def extra_repr(self) -> str:
return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"
def flops(self):
flops = 0
for blk in self.blocks:
flops += blk.flops()
if self.upsample is not None:
flops += self.upsample.flops()
return flopsStyleSwinTransformerBlock():
class StyleSwinTransformerBlock(nn.Module): # Swin block
r""" StyleSwin Transformer Block.
Args:
dim (int): Number of input channels.
input_resolution (tuple[int]): Input resulotion.
num_heads (int): Number of attention heads.
window_size (int): Window size.
shift_size (int): Shift size for SW-MSA.
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
drop (float, optional): Dropout rate. Default: 0.0
attn_drop (float, optional): Attention dropout rate. Default: 0.0
drop_path (float, optional): Stochastic depth rate. Default: 0.0
act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
style_dim (int): Dimension of style vector.
"""输入的参数：和 StyleBasicLayer()大部分的参数相似；参考StyleBasicLayer()
dim (int): 输入的通道数，注意这边输入的是嵌入后的，而不是最早的图像通道数
input_resolution (tuple[int]): 宽度高度方向分别分成了几个patch
num_heads (int): 注意力的head数量
window_size (int):窗口的尺寸
shift_size (int):SW-MSA滑移窗口每次的移动距离
mlp_ratio (float): MLP的隐藏层维度与patch embedding 维度的比例
qkv_bias (bool, optional):为True时，则给QKV分别加一个可学习的偏置项
qk_scale (float | None, optional): q的放缩因子，若未设置则使用head_dim ** -0.5
drop (float, optional): Dropout比例，默认为0，用于防止过拟合
attn_drop (float, optional): Attention dropout比例，默认为0
drop_path (float, optional): Stochastic depth rate. Default: 0.0
act_layer (nn.Module, optional): 激活函数层，默认使用nn.GELU
norm_layer (nn.Module, optional): 归一化层，默认为nn.LayerNorm
def __init__(self, dim, input_resolution, num_heads, window_size=7,
mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
act_layer=nn.GELU, style_dim=512):
super().__init__()
self.dim = dim
self.input_resolution = input_resolution
self.num_heads = num_heads
self.window_size = window_size
self.mlp_ratio = mlp_ratio
self.shift_size = self.window_size // 2
self.style_dim = style_dim
if min(self.input_resolution) <= self.window_size:
# if window size is larger than input resolution, we don't partition windows
self.shift_size = 0
self.window_size = min(self.input_resolution)
assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
self.norm1 = AdaptiveInstanceNorm(dim, style_dim)
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.proj = nn.Linear(dim, dim)
self.attn = nn.ModuleList([
WindowAttention(
dim // 2, window_size=to_2tuple(self.window_size), num_heads=num_heads // 2,
qk_scale=qk_scale, attn_drop=attn_drop),
WindowAttention(
dim // 2, window_size=to_2tuple(self.window_size), num_heads=num_heads // 2,
qk_scale=qk_scale, attn_drop=attn_drop),
])
attn_mask1 = None
attn_mask2 = None
if self.shift_size > 0:
# calculate attention mask for SW-MSA
H, W = self.input_resolution
img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1
h_slices = (slice(0, -self.window_size),
slice(-self.window_size, -self.shift_size),
slice(-self.shift_size, None))
w_slices = (slice(0, -self.window_size),
slice(-self.window_size, -self.shift_size),
slice(-self.shift_size, None))
cnt = 0
for h in h_slices:
for w in w_slices:
img_mask[:, h, w, :] = cnt
cnt += 1
# nW, window_size, window_size, 1
mask_windows = window_partition(img_mask, self.window_size)
mask_windows = mask_windows.view(-1,
self.window_size * self.window_size)
attn_mask2 = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
attn_mask2 = attn_mask2.masked_fill(
attn_mask2 != 0, float(-100.0)).masked_fill(attn_mask2 == 0, float(0.0))
self.register_buffer("attn_mask1", attn_mask1)
self.register_buffer("attn_mask2", attn_mask2)
self.norm2 = AdaptiveInstanceNorm(dim, style_dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) def forward(self, x, style):
H, W = self.input_resolution # 解析输入图像的分辨率
B, L, C = x.shape # 解析输入的维度
assert L == H * W, "input feature has wrong size" # 判断L和H*W是否一致
# Double Attn
shortcut = x # 进行原始身份的映射
x = self.norm1(x.transpose(-1, -2), style).transpose(-1, -2)
qkv = self.qkv(x).reshape(B, -1, 3, C).permute(2, 0, 1, 3).reshape(3 * B, H, W, C)
qkv_1 = qkv[:, :, :, : C // 2].reshape(3, B, H, W, C // 2)
if self.shift_size > 0: # 判断是否进行循环移位
qkv_2 = torch.roll(qkv[:, :, :, C // 2:], shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)).reshape(3, B, H, W, C // 2)
else:
qkv_2 = qkv[:, :, :, C // 2:].reshape(3, B, H, W, C // 2)
q1_windows, k1_windows, v1_windows = self.get_window_qkv(qkv_1)
q2_windows, k2_windows, v2_windows = self.get_window_qkv(qkv_2)
x1 = self.attn[0](q1_windows, k1_windows, v1_windows, self.attn_mask1)
x2 = self.attn[1](q2_windows, k2_windows, v2_windows, self.attn_mask2)
x1 = window_reverse(x1.view(-1, self.window_size * self.window_size, C // 2), self.window_size, H, W)
x2 = window_reverse(x2.view(-1, self.window_size * self.window_size, C // 2), self.window_size, H, W)
if self.shift_size > 0: # 循环移位进行恢复
x2 = torch.roll(x2, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
else:
x2 = x2
x = torch.cat([x1.reshape(B, H * W, C // 2), x2.reshape(B, H * W, C // 2)], dim=2)
x = self.proj(x)
# FFN
x = shortcut + x
x = x + self.mlp(self.norm2(x.transpose(-1, -2), style).transpose(-1, -2))
return x分别计算q,k,v
def get_window_qkv(self, qkv):
q, k, v = qkv[0], qkv[1], qkv[2] # B, H, W, C
C = q.shape[-1]
q_windows = window_partition(q, self.window_size).view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C
k_windows = window_partition(k, self.window_size).view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C
v_windows = window_partition(v, self.window_size).view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C
return q_windows, k_windows, v_windows def extra_repr(self) -> str:
return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \
f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
def flops(self):
flops = 0
H, W = self.input_resolution
# norm1
flops += 1 * self.style_dim * self.dim * 2
flops += 2 * (H * W) * self.dim
# W-MSA/SW-MSA
nW = H * W / self.window_size / self.window_size
for attn in self.attn:
flops += nW * (attn.flops(self.window_size * self.window_size))
# mlp
flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio
# norm2
flops += 1 * self.style_dim * self.dim * 2
flops += 2 * (H * W) * self.dim
return flopsBilinearUpsample():
class BilinearUpsample(nn.Module):
""" BilinearUpsample Layer.
Args:
input_resolution (tuple[int]): Resolution of input feature.
dim (int): Number of input channels.
out_dim (int): Number of output channels.
"""
def __init__(self, input_resolution, dim, out_dim=None):
super().__init__()
assert dim % 2 == 0, f"x dim are not even."
self.upsample = nn.Upsample(scale_factor=2, mode='bilinear')
self.norm = nn.LayerNorm(dim)
self.reduction = nn.Linear(dim, out_dim, bias=False)
self.input_resolution = input_resolution
self.dim = dim
self.out_dim = out_dim
self.alpha = nn.Parameter(torch.zeros(1))
self.sin_pos_embed = SinusoidalPositionalEmbedding(embedding_dim=out_dim // 2, padding_idx=0, init_size=out_dim // 2)
def forward(self, x):
"""
x: B, H*W, C
"""
H, W = self.input_resolution
B, L, C = x.shape
assert L == H * W, "input feature has wrong size"
assert C == self.dim, "wrong in PatchMerging"
x = x.view(B, H, W, -1)
x = x.permute(0, 3, 1, 2).contiguous() # B,C,H,W
x = self.upsample(x)
x = x.permute(0, 2, 3, 1).contiguous().view(B, L*4, C) # B,H,W,C
x = self.norm(x)
x = self.reduction(x)
# Add SPE
x = x.reshape(B, H * 2, W * 2, self.out_dim).permute(0, 3, 1, 2)
x += self.sin_pos_embed.make_grid2d(H * 2, W * 2, B) * self.alpha
x = x.permute(0, 2, 3, 1).contiguous().view(B, H * 2 * W * 2, self.out_dim)
return x
def extra_repr(self) -> str:
return f"input_resolution={self.input_resolution}, dim={self.dim}"
def flops(self):
H, W = self.input_resolution
# LN
flops = 4 * H * W * self.dim
# proj
flops += 4 * H * W * self.dim * (self.out_dim)
# SPE
flops += 4 * H * W * 2
# bilinear
flops += 4 * self.input_resolution[0] * self.input_resolution[1] * self.dim * 5
return flopsConstantInput():
class ConstantInput(nn.Module):
def __init__(self, channel, size=4):
super().__init__()
self.input = nn.Parameter(torch.randn(1, channel, size, size))
def forward(self, input):
batch = input.shape[0]
out = self.input.repeat(batch, 1, 1, 1)
return out # return一个正态分布采样、参数化的tensor，维度为（batchsize，512，4，4）常量初始化为1，是一个可学习的参数，要参与训练，选择用Parameter进行封装。
Generator(): self.size是生成图像大小，self.style_dim是隐层码维度，由输入参数指定。
class Generator(nn.Module):
def __init__(
self,
size, # 生成图像大小
style_dim, # 隐层码维度=512
n_mlp, # z到w的映射网络层数
channel_multiplier=2,
lr_mlp=0.01,
enable_full_resolution=8,
mlp_ratio=4,
use_checkpoint=False,
qkv_bias=True,
qk_scale=None,
drop_rate=0,
attn_drop_rate=0,
):
super().__init__()
self.style_dim = style_dim
self.size = size
self.mlp_ratio = mlp_ratioself.style是由PixelNorm层和8个EqualLinear层组成的MLP，也就是将噪声z映射为隐层码w的网络。
layers = [PixelNorm()]
for _ in range(n_mlp): # n_mlp=8 style_dim=512 lr_mlp=0.01
layers.append(
EqualLinear(
style_dim, style_dim, lr_mul=lr_mlp, activation='fused_lrelu'
)
)
self.style = nn.Sequential(*layers) # self.style是由PixelNorm层和8个EqualLinear层组成的MLP，也就是将噪声z映射为隐层码w的网络。start = 2
depths = [2, 2, 2, 2, 2, 2, 2, 2, 2]
in_channels = [
512, # 4
512, # 8
512, # 16
512, # 32
256 * channel_multiplier, # 64 channel_multiplier=2
128 * channel_multiplier, # 128
64 * channel_multiplier, # 256
32 * channel_multiplier, # 512
16 * channel_multiplier # 1024
]
end = int(math.log(size, 2))
num_heads = [max(c // 32, 4) for c in in_channels]
full_resolution_index = int(math.log(enable_full_resolution, 2))
window_sizes = [2 ** i if i <= full_resolution_index else 8 for i in range(start, end + 1)]
self.input = ConstantInput(in_channels[0]) # self.input是主支的常量输入 in_channels[0]=512
self.layers = nn.ModuleList()
self.to_rgbs = nn.ModuleList()
num_layers = 0对StyleBasicLayer进行循环
for i_layer in range(start, end + 1):
in_channel = in_channels[i_layer - start]
layer = StyleBasicLayer(dim=in_channel,
input_resolution=(2 ** i_layer,2 ** i_layer),
depth=depths[i_layer - start],
num_heads=num_heads[i_layer - start],
window_size=window_sizes[i_layer - start],
out_dim=in_channels[i_layer - start + 1] if (i_layer < end) else None,
mlp_ratio=self.mlp_ratio,
qkv_bias=qkv_bias, qk_scale=qk_scale,
drop=drop_rate, attn_drop=attn_drop_rate,
upsample=BilinearUpsample if (i_layer < end) else None,
use_checkpoint=use_checkpoint, style_dim=style_dim)
self.layers.append(layer)
out_dim = in_channels[i_layer - start + 1] if (i_layer < end) else in_channels[i_layer - start]
upsample = True if (i_layer < end) else False
to_rgb = ToRGB(out_dim, upsample=upsample, resolution=(2 ** i_layer))
self.to_rgbs.append(to_rgb)
num_layers += 2
self.n_latent = num_layers
self.apply(self._init_weights)_init_weights(): 权重，用权重来控制样式
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
if m.bias is not None:
nn.init.constant_(m.bias, 0)
if m.weight is not None:
nn.init.constant_(m.weight, 1.0)
elif isinstance(m, nn.Conv2d):
nn.init.xavier_normal_(m.weight, gain=.02)
if hasattr(m, 'bias') and m.bias is not None:
nn.init.constant_(m.bias, 0)
def forward(
self,
noise,
return_latents=False,
inject_index=None,
truncation=1,
truncation_latent=None,
):
styles = self.style(noise)
inject_index = self.n_latent
if truncation < 1:
style_t = []
for style in styles:
style_t.append(
truncation_latent + truncation * (style - truncation_latent)
)
styles = torch.cat(style_t, dim=0)
if styles.ndim < 3:
latent = styles.unsqueeze(1).repeat(1, inject_index, 1)
else:
latent = styles
x = self.input(latent)
B, C, H, W = x.shape
x = x.permute(0, 2, 3, 1).contiguous().view(B, H * W, C)
count = 0
skip = None
for layer, to_rgb in zip(self.layers, self.to_rgbs):
x = layer(x, latent[:,count,:], latent[:,count+1,:])
b, n, c = x.shape
h, w = int(math.sqrt(n)), int(math.sqrt(n))
skip = to_rgb(x.transpose(-1, -2).reshape(b, c, h, w), skip)
count = count + 2
B, L, C = x.shape
assert L == self.size * self.size
x = x.reshape(B, self.size, self.size, C).permute(0, 3, 1, 2).contiguous()
image = skip
if return_latents:
return image, latent
else:
return image, None
def flops(self):
flops = 0
for _, layer in enumerate(self.layers):
flops += layer.flops()
for _, layer in enumerate(self.to_rgbs):
flops += layer.flops()
# 8 FC + PixelNorm
flops += 1 * 10 * self.style_dim * self.style_dim
return flops以上有理解不对的地方，欢迎讨论批评指正~~