126 lines
5.2 KiB
Python
126 lines
5.2 KiB
Python
import math
|
|
from typing import Optional, List
|
|
|
|
import torch
|
|
from torch import Tensor
|
|
from torch.nn import Module, ModuleList, Sequential, Upsample
|
|
|
|
from tha3.nn.common.conv_block_factory import ConvBlockFactory
|
|
from tha3.nn.nonlinearity_factory import LeakyReLUFactory
|
|
from tha3.nn.normalization import InstanceNorm2dFactory
|
|
from tha3.nn.util import BlockArgs
|
|
|
|
|
|
class ResizeConvEncoderDecoderArgs:
|
|
def __init__(self,
|
|
image_size: int,
|
|
input_channels: int,
|
|
start_channels: int,
|
|
bottleneck_image_size,
|
|
num_bottleneck_blocks,
|
|
max_channels: int,
|
|
block_args: Optional[BlockArgs] = None,
|
|
upsample_mode: str = 'bilinear',
|
|
use_separable_convolution=False):
|
|
self.use_separable_convolution = use_separable_convolution
|
|
self.upsample_mode = upsample_mode
|
|
self.block_args = block_args
|
|
self.max_channels = max_channels
|
|
self.num_bottleneck_blocks = num_bottleneck_blocks
|
|
self.bottleneck_image_size = bottleneck_image_size
|
|
self.start_channels = start_channels
|
|
self.image_size = image_size
|
|
self.input_channels = input_channels
|
|
|
|
|
|
class ResizeConvEncoderDecoder(Module):
|
|
def __init__(self, args: ResizeConvEncoderDecoderArgs):
|
|
super().__init__()
|
|
self.args = args
|
|
|
|
self.num_levels = int(math.log2(args.image_size // args.bottleneck_image_size)) + 1
|
|
|
|
conv_block_factory = ConvBlockFactory(args.block_args, args.use_separable_convolution)
|
|
|
|
self.downsample_blocks = ModuleList()
|
|
self.downsample_blocks.append(conv_block_factory.create_conv7_block(args.input_channels, args.start_channels))
|
|
current_image_size = args.image_size
|
|
current_num_channels = args.start_channels
|
|
while current_image_size > args.bottleneck_image_size:
|
|
next_image_size = current_image_size // 2
|
|
next_num_channels = self.get_num_output_channels_from_image_size(next_image_size)
|
|
self.downsample_blocks.append(conv_block_factory.create_downsample_block(
|
|
in_channels=current_num_channels,
|
|
out_channels=next_num_channels,
|
|
is_output_1x1=False))
|
|
current_image_size = next_image_size
|
|
current_num_channels = next_num_channels
|
|
assert len(self.downsample_blocks) == self.num_levels
|
|
|
|
self.bottleneck_blocks = ModuleList()
|
|
for i in range(args.num_bottleneck_blocks):
|
|
self.bottleneck_blocks.append(conv_block_factory.create_resnet_block(current_num_channels, is_1x1=False))
|
|
|
|
self.output_image_sizes = [current_image_size]
|
|
self.output_num_channels = [current_num_channels]
|
|
self.upsample_blocks = ModuleList()
|
|
if args.upsample_mode == 'nearest':
|
|
align_corners = None
|
|
else:
|
|
align_corners = False
|
|
while current_image_size < args.image_size:
|
|
next_image_size = current_image_size * 2
|
|
next_num_channels = self.get_num_output_channels_from_image_size(next_image_size)
|
|
self.upsample_blocks.append(
|
|
Sequential(
|
|
Upsample(scale_factor=2, mode=args.upsample_mode, align_corners=align_corners),
|
|
conv_block_factory.create_conv3_block(
|
|
in_channels=current_num_channels, out_channels=next_num_channels)))
|
|
current_image_size = next_image_size
|
|
current_num_channels = next_num_channels
|
|
self.output_image_sizes.append(current_image_size)
|
|
self.output_num_channels.append(current_num_channels)
|
|
|
|
def get_num_output_channels_from_level(self, level: int):
|
|
return self.get_num_output_channels_from_image_size(self.args.image_size // (2 ** level))
|
|
|
|
def get_num_output_channels_from_image_size(self, image_size: int):
|
|
return min(self.args.start_channels * (self.args.image_size // image_size), self.args.max_channels)
|
|
|
|
def forward(self, feature: Tensor) -> List[Tensor]:
|
|
outputs = []
|
|
for block in self.downsample_blocks:
|
|
feature = block(feature)
|
|
for block in self.bottleneck_blocks:
|
|
feature = block(feature)
|
|
outputs.append(feature)
|
|
for block in self.upsample_blocks:
|
|
feature = block(feature)
|
|
outputs.append(feature)
|
|
return outputs
|
|
|
|
|
|
if __name__ == "__main__":
|
|
device = torch.device('cuda')
|
|
args = ResizeConvEncoderDecoderArgs(
|
|
image_size=512,
|
|
input_channels=4 + 6,
|
|
start_channels=32,
|
|
bottleneck_image_size=32,
|
|
num_bottleneck_blocks=6,
|
|
max_channels=512,
|
|
use_separable_convolution=True,
|
|
block_args=BlockArgs(
|
|
initialization_method='he',
|
|
use_spectral_norm=False,
|
|
normalization_layer_factory=InstanceNorm2dFactory(),
|
|
nonlinearity_factory=LeakyReLUFactory(inplace=False, negative_slope=0.1)))
|
|
module = ResizeConvEncoderDecoder(args).to(device)
|
|
print(module.output_image_sizes)
|
|
print(module.output_num_channels)
|
|
|
|
input = torch.zeros(8, 4 + 6, 512, 512, device=device)
|
|
outputs = module(input)
|
|
for output in outputs:
|
|
print(output.shape)
|