Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,814 Bytes
3e648fb |
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 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
class LayerNorm(nn.Module):
r""" LayerNorm that supports two data formats: channels_last (default) or channels_first.
The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
shape (batch_size, height, width, channels) while channels_first corresponds to inputs
with shape (batch_size, channels, height, width).
"""
def __init__(self, normalized_shape, eps=1e-6, data_format="channels_first"):
super().__init__()
self.weight = nn.Parameter(torch.ones(normalized_shape))
self.bias = nn.Parameter(torch.zeros(normalized_shape))
self.eps = eps
self.data_format = data_format
if self.data_format not in ["channels_last", "channels_first"]:
raise NotImplementedError
self.normalized_shape = (normalized_shape, )
def forward(self, x):
if self.data_format == "channels_last":
return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
elif self.data_format == "channels_first":
u = x.mean(1, keepdim=True)
s = (x - u).pow(2).mean(1, keepdim=True)
x = (x - u) / torch.sqrt(s + self.eps)
x = self.weight[:, None, None] * x + self.bias[:, None, None]
return x
class NormDownsample(nn.Module):
def __init__(self,in_ch,out_ch,scale=0.5,use_norm=False):
super(NormDownsample, self).__init__()
self.use_norm=use_norm
if self.use_norm:
self.norm=LayerNorm(out_ch)
self.prelu = nn.PReLU()
self.down = nn.Sequential(
nn.Conv2d(in_ch, out_ch,kernel_size=3,stride=1, padding=1, bias=False),
nn.UpsamplingBilinear2d(scale_factor=scale))
def forward(self, x):
x = self.down(x)
x = self.prelu(x)
if self.use_norm:
x = self.norm(x)
return x
else:
return x
class NormUpsample(nn.Module):
def __init__(self, in_ch,out_ch,scale=2,use_norm=False):
super(NormUpsample, self).__init__()
self.use_norm=use_norm
if self.use_norm:
self.norm=LayerNorm(out_ch)
self.prelu = nn.PReLU()
self.up_scale = nn.Sequential(
nn.Conv2d(in_ch,out_ch,kernel_size=3,stride=1, padding=1, bias=False),
nn.UpsamplingBilinear2d(scale_factor=scale))
self.up = nn.Conv2d(out_ch*2,out_ch,kernel_size=1,stride=1, padding=0, bias=False)
def forward(self, x,y):
x = self.up_scale(x)
x = torch.cat([x, y],dim=1)
x = self.up(x)
x = self.prelu(x)
if self.use_norm:
return self.norm(x)
else:
return x
|