|
import torch.nn.functional as F |
|
from torch import nn |
|
|
|
|
|
class PreactResBlock(nn.Sequential): |
|
def __init__(self, dim): |
|
super().__init__( |
|
nn.GroupNorm(dim // 16, dim), |
|
nn.GELU(), |
|
nn.Conv2d(dim, dim, 3, padding=1), |
|
nn.GroupNorm(dim // 16, dim), |
|
nn.GELU(), |
|
nn.Conv2d(dim, dim, 3, padding=1), |
|
) |
|
|
|
def forward(self, x): |
|
return x + super().forward(x) |
|
|
|
|
|
class UNetBlock(nn.Module): |
|
def __init__(self, input_dim, output_dim=None, scale_factor=1.0): |
|
super().__init__() |
|
if output_dim is None: |
|
output_dim = input_dim |
|
self.pre_conv = nn.Conv2d(input_dim, output_dim, 3, padding=1) |
|
self.res_block1 = PreactResBlock(output_dim) |
|
self.res_block2 = PreactResBlock(output_dim) |
|
self.downsample = self.upsample = nn.Identity() |
|
if scale_factor > 1: |
|
self.upsample = nn.Upsample(scale_factor=scale_factor) |
|
elif scale_factor < 1: |
|
self.downsample = nn.Upsample(scale_factor=scale_factor) |
|
|
|
def forward(self, x, h=None): |
|
""" |
|
Args: |
|
x: (b c h w), last output |
|
h: (b c h w), skip output |
|
Returns: |
|
o: (b c h w), output |
|
s: (b c h w), skip output |
|
""" |
|
x = self.upsample(x) |
|
if h is not None: |
|
assert x.shape == h.shape, f"{x.shape} != {h.shape}" |
|
x = x + h |
|
x = self.pre_conv(x) |
|
x = self.res_block1(x) |
|
x = self.res_block2(x) |
|
return self.downsample(x), x |
|
|
|
|
|
class UNet(nn.Module): |
|
def __init__( |
|
self, input_dim, output_dim, hidden_dim=16, num_blocks=4, num_middle_blocks=2 |
|
): |
|
super().__init__() |
|
self.input_dim = input_dim |
|
self.output_dim = output_dim |
|
self.input_proj = nn.Conv2d(input_dim, hidden_dim, 3, padding=1) |
|
self.encoder_blocks = nn.ModuleList( |
|
[ |
|
UNetBlock( |
|
input_dim=hidden_dim * 2**i, |
|
output_dim=hidden_dim * 2 ** (i + 1), |
|
scale_factor=0.5, |
|
) |
|
for i in range(num_blocks) |
|
] |
|
) |
|
self.middle_blocks = nn.ModuleList( |
|
[ |
|
UNetBlock(input_dim=hidden_dim * 2**num_blocks) |
|
for _ in range(num_middle_blocks) |
|
] |
|
) |
|
self.decoder_blocks = nn.ModuleList( |
|
[ |
|
UNetBlock( |
|
input_dim=hidden_dim * 2 ** (i + 1), |
|
output_dim=hidden_dim * 2**i, |
|
scale_factor=2, |
|
) |
|
for i in reversed(range(num_blocks)) |
|
] |
|
) |
|
self.head = nn.Sequential( |
|
nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1), |
|
nn.GELU(), |
|
nn.Conv2d(hidden_dim, output_dim, 1), |
|
) |
|
|
|
@property |
|
def scale_factor(self): |
|
return 2 ** len(self.encoder_blocks) |
|
|
|
def pad_to_fit(self, x): |
|
""" |
|
Args: |
|
x: (b c h w), input |
|
Returns: |
|
x: (b c h' w'), padded input |
|
""" |
|
hpad = (self.scale_factor - x.shape[2] % self.scale_factor) % self.scale_factor |
|
wpad = (self.scale_factor - x.shape[3] % self.scale_factor) % self.scale_factor |
|
return F.pad(x, (0, wpad, 0, hpad)) |
|
|
|
def forward(self, x): |
|
""" |
|
Args: |
|
x: (b c h w), input |
|
Returns: |
|
o: (b c h w), output |
|
""" |
|
shape = x.shape |
|
|
|
x = self.pad_to_fit(x) |
|
x = self.input_proj(x) |
|
|
|
s_list = [] |
|
for block in self.encoder_blocks: |
|
x, s = block(x) |
|
s_list.append(s) |
|
|
|
for block in self.middle_blocks: |
|
x, _ = block(x) |
|
|
|
for block, s in zip(self.decoder_blocks, reversed(s_list)): |
|
x, _ = block(x, s) |
|
|
|
x = self.head(x) |
|
x = x[..., : shape[2], : shape[3]] |
|
|
|
return x |
|
|
|
def test(self, shape=(3, 512, 256)): |
|
import ptflops |
|
|
|
macs, params = ptflops.get_model_complexity_info( |
|
self, |
|
shape, |
|
as_strings=True, |
|
print_per_layer_stat=True, |
|
verbose=True, |
|
) |
|
|
|
print(f"macs: {macs}") |
|
print(f"params: {params}") |
|
|
|
|
|
def main(): |
|
model = UNet(3, 3) |
|
model.test() |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
