import math
import torch
from torch import nn
import torch.nn.functional as F
from inspect import isfunction
from kornia.filters import gaussian_blur2d
def exists(x):
return x is not None
def default(val, d):
if exists(val):
return val
return d() if isfunction(d) else d
class PositionalEncoding(nn.Module):
def __init__(self, dim):
super().__init__()
self.dim = dim
def forward(self, noise_level):
count = self.dim // 2
step = torch.arange(count, dtype=noise_level.dtype,
device=noise_level.device) / count
encoding = noise_level.unsqueeze(
1) * torch.exp(-math.log(1e4) * step.unsqueeze(0))
encoding = torch.cat(
[torch.sin(encoding), torch.cos(encoding)], dim=-1)
return encoding
class FeatureWiseAffine(nn.Module):
def __init__(self, in_channels, out_channels, use_affine_level=False):
super(FeatureWiseAffine, self).__init__()
self.use_affine_level = use_affine_level
self.noise_func = nn.Sequential(
nn.Linear(in_channels, out_channels * (1 + self.use_affine_level))
)
def forward(self, x, noise_embed):
batch = x.shape[0]
if self.use_affine_level:
gamma, beta = self.noise_func(noise_embed).view(
batch, -1, 1, 1).chunk(2, dim=1)
x = (1 + gamma) * x + beta
else:
x = x + self.noise_func(noise_embed).view(batch, -1, 1, 1)
return x
class Swish(nn.Module):
def forward(self, x):
return x * torch.sigmoid(x)
class Upsample(nn.Module):
def __init__(self, dim):
super().__init__()
self.up = nn.Upsample(scale_factor=2, mode="nearest")
self.conv = nn.Conv2d(dim, dim, 3, padding=1)
def forward(self, x):
return self.conv(self.up(x))
class Downsample(nn.Module):
def __init__(self, dim):
super().__init__()
self.conv = nn.Conv2d(dim, dim, 3, 2, 1)
def forward(self, x):
return self.conv(x)
class Block(nn.Module):
def __init__(self, dim, dim_out, groups=32, dropout=0):
super().__init__()
self.block = nn.Sequential(
nn.GroupNorm(groups, dim),
Swish(),
nn.Dropout(dropout) if dropout != 0 else nn.Identity(),
nn.Conv2d(dim, dim_out, 3, padding=1)
)
def forward(self, x):
return self.block(x)
class ResnetBlock(nn.Module):
def __init__(self, dim, dim_out, noise_level_emb_dim=None, dropout=0, use_affine_level=False, norm_groups=32):
super().__init__()
self.noise_func = FeatureWiseAffine(
noise_level_emb_dim, dim_out, use_affine_level)
self.block1 = Block(dim, dim_out, groups=norm_groups)
self.block2 = Block(dim_out, dim_out, groups=norm_groups, dropout=dropout)
self.res_conv = nn.Conv2d(
dim, dim_out, 1) if dim != dim_out else nn.Identity()
def forward(self, x, time_emb):
b, c, h, w = x.shape
h = self.block1(x)
h = self.noise_func(h, time_emb)
h = self.block2(h)
return h + self.res_conv(x)
class SelfAttention(nn.Module):
def __init__(self, in_channel, n_head=1, norm_groups=32):
super().__init__()
self.n_head = n_head
self.norm = nn.GroupNorm(norm_groups, in_channel)
self.qkv = nn.Conv2d(in_channel, in_channel * 3, 1, bias=False)
self.out = nn.Conv2d(in_channel, in_channel, 1)
def forward(self, input):
batch, channel, height, width = input.shape
n_head = self.n_head
head_dim = channel // n_head
norm = self.norm(input)
qkv = self.qkv(norm).view(batch, n_head, head_dim * 3, height, width)
query, key, value = qkv.chunk(3, dim=2) # bhdyx
attn = torch.einsum(
"bnchw, bncyx -> bnhwyx", query, key
).contiguous() / math.sqrt(channel)
attn = attn.view(batch, n_head, height, width, -1)
attn = torch.softmax(attn, -1)
attn = attn.view(batch, n_head, height, width, height, width)
out = torch.einsum("bnhwyx, bncyx -> bnchw", attn, value).contiguous()
out = self.out(out.view(batch, channel, height, width))
return out + input
class ResnetBlocWithAttn(nn.Module):
def __init__(self, dim, dim_out, *, noise_level_emb_dim=None, norm_groups=32, dropout=0, with_attn=False):
super().__init__()
self.with_attn = with_attn
self.res_block = ResnetBlock(
dim, dim_out, noise_level_emb_dim, norm_groups=norm_groups, dropout=dropout)
if with_attn:
self.attn = SelfAttention(dim_out, norm_groups=norm_groups)
def forward(self, x, time_emb):
x = self.res_block(x, time_emb)
if (self.with_attn):
x = self.attn(x)
return x
class FCB(nn.Module):
def __init__(self, channel, kernel_size=3):
super().__init__()
self.ks = kernel_size
self.sigma_rate = 1
params = torch.ones((4, 1), requires_grad=True)
self.params = nn.Parameter(params)
def forward(self, x):
#
x1 = gaussian_blur2d(x, (self.ks, self.ks), (1 * self.sigma_rate, 1 * self.sigma_rate))
R1 = x - x1
x2 = gaussian_blur2d(x, (self.ks * 2 - 1, self.ks * 2 - 1), (2 * self.sigma_rate, 2 * self.sigma_rate))
x3 = gaussian_blur2d(x, (self.ks * 4 - 1, self.ks * 4 - 1), (4 * self.sigma_rate, 4 * self.sigma_rate))
R2 = x1 - x2
R3 = x2 - x3
R1 = R1.unsqueeze(dim=-1)
R2 = R2.unsqueeze(dim=-1)
R3 = R3.unsqueeze(dim=-1)
R_cat = torch.cat([R1, R2, R3, x.unsqueeze(dim=-1)], dim=-1)
sum_ = torch.matmul(R_cat, self.params).squeeze(dim=-1)
return sum_
class UNet(nn.Module):
def __init__(
self,
in_channel=6,
out_channel=3,
inner_channel=32,
norm_groups=32,
channel_mults=(1, 2, 4, 8, 8),
attn_res=[8],
res_blocks=3,
dropout=0,
with_noise_level_emb=True,
image_size=128,
fcb=True
):
super().__init__()
self.fcb = fcb
if with_noise_level_emb:
noise_level_channel = inner_channel
self.noise_level_mlp = nn.Sequential(
PositionalEncoding(inner_channel),
nn.Linear(inner_channel, inner_channel * 4),
Swish(),
nn.Linear(inner_channel * 4, inner_channel)
)
else:
noise_level_channel = None
self.noise_level_mlp = None
num_mults = len(channel_mults)
pre_channel = inner_channel
feat_channels = [pre_channel]
now_res = image_size
downs = [nn.Conv2d(in_channel, inner_channel,
kernel_size=3, padding=1)]
for ind in range(num_mults):
is_last = (ind == num_mults - 1)
use_attn = (now_res in attn_res)
channel_mult = inner_channel * channel_mults[ind]
for _ in range(0, res_blocks):
downs.append(ResnetBlocWithAttn(
pre_channel, channel_mult, noise_level_emb_dim=noise_level_channel, norm_groups=norm_groups,
dropout=dropout, with_attn=use_attn))
feat_channels.append(channel_mult)
pre_channel = channel_mult
if not is_last:
downs.append(Downsample(pre_channel))
feat_channels.append(pre_channel)
now_res = now_res // 2
self.downs = nn.ModuleList(downs)
self.mid = nn.ModuleList([
ResnetBlocWithAttn(pre_channel, pre_channel, noise_level_emb_dim=noise_level_channel,
norm_groups=norm_groups,
dropout=dropout, with_attn=True),
ResnetBlocWithAttn(pre_channel, pre_channel, noise_level_emb_dim=noise_level_channel,
norm_groups=norm_groups,
dropout=dropout, with_attn=False)
])
ups = []
fbs = []
for ind in reversed(range(num_mults)):
is_last = (ind < 1)
use_attn = (now_res in attn_res)
channel_mult = inner_channel * channel_mults[ind]
for _ in range(0, res_blocks + 1):
ups.append(ResnetBlocWithAttn(
pre_channel + feat_channels.pop(), channel_mult, noise_level_emb_dim=noise_level_channel,
norm_groups=norm_groups,
dropout=dropout, with_attn=use_attn))
pre_channel = channel_mult
tmp = FCB(pre_channel) if self.fcb else pre_channel
fbs.append(tmp)
if not is_last:
ups.append(Upsample(pre_channel))
tmp = FCB(pre_channel) if self.fcb else pre_channel
fbs.append(tmp)
now_res = now_res * 2
self.ups = nn.ModuleList(ups)
self.fbs = nn.ModuleList(fbs)
self.final_conv = Block(pre_channel, default(out_channel, in_channel), groups=norm_groups)
def forward(self, x, time):
t = self.noise_level_mlp(time) if exists(
self.noise_level_mlp) else None
feats = []
for layer in self.downs:
if isinstance(layer, ResnetBlocWithAttn):
x = layer(x, t)
else:
x = layer(x)
feats.append(x)
for layer in self.mid:
if isinstance(layer, ResnetBlocWithAttn):
x = layer(x, t)
else:
x = layer(x)
for layer, fb in zip(self.ups, self.fbs):
if isinstance(layer, ResnetBlocWithAttn):
tmp = feats.pop()
if self.fcb:
tmp = fb(tmp)
x = layer(torch.cat((x, tmp), dim=1), t)
else:
x = layer(x)
tmp = self.final_conv(x)
return tmp
标签：11,dim,noise,nn,level,self,channel From： https://www.cnblogs.com/yyhappy/p/17929190.html