| from typing import Optional, List | |
| import torch | |
| from torch import Tensor | |
| from torch.nn import ModuleList, Module, Upsample | |
| from tha3.nn.common.conv_block_factory import ConvBlockFactory | |
| from tha3.nn.nonlinearity_factory import ReLUFactory | |
| from tha3.nn.normalization import InstanceNorm2dFactory | |
| from tha3.nn.util import BlockArgs | |
| class ResizeConvUNetArgs: | |
| def __init__(self, | |
| image_size: int, | |
| input_channels: int, | |
| start_channels: int, | |
| bottleneck_image_size: int, | |
| num_bottleneck_blocks: int, | |
| max_channels: int, | |
| upsample_mode: str = 'bilinear', | |
| block_args: Optional[BlockArgs] = None, | |
| use_separable_convolution: bool = False): | |
| if block_args is None: | |
| block_args = BlockArgs( | |
| normalization_layer_factory=InstanceNorm2dFactory(), | |
| nonlinearity_factory=ReLUFactory(inplace=False)) | |
| self.use_separable_convolution = use_separable_convolution | |
| self.block_args = block_args | |
| self.upsample_mode = upsample_mode | |
| self.max_channels = max_channels | |
| self.num_bottleneck_blocks = num_bottleneck_blocks | |
| self.bottleneck_image_size = bottleneck_image_size | |
| self.input_channels = input_channels | |
| self.start_channels = start_channels | |
| self.image_size = image_size | |
| class ResizeConvUNet(Module): | |
| def __init__(self, args: ResizeConvUNetArgs): | |
| super().__init__() | |
| self.args = args | |
| conv_block_factory = ConvBlockFactory(args.block_args, args.use_separable_convolution) | |
| self.downsample_blocks = ModuleList() | |
| self.downsample_blocks.append(conv_block_factory.create_conv3_block( | |
| self.args.input_channels, | |
| self.args.start_channels)) | |
| current_channels = self.args.start_channels | |
| current_size = self.args.image_size | |
| size_to_channel = { | |
| current_size: current_channels | |
| } | |
| while current_size > self.args.bottleneck_image_size: | |
| next_size = current_size // 2 | |
| next_channels = min(self.args.max_channels, current_channels * 2) | |
| self.downsample_blocks.append(conv_block_factory.create_downsample_block( | |
| current_channels, | |
| next_channels, | |
| is_output_1x1=False)) | |
| current_size = next_size | |
| current_channels = next_channels | |
| size_to_channel[current_size] = current_channels | |
| self.bottleneck_blocks = ModuleList() | |
| for i in range(self.args.num_bottleneck_blocks): | |
| self.bottleneck_blocks.append(conv_block_factory.create_resnet_block(current_channels, is_1x1=False)) | |
| self.output_image_sizes = [current_size] | |
| self.output_num_channels = [current_channels] | |
| self.upsample_blocks = ModuleList() | |
| while current_size < self.args.image_size: | |
| next_size = current_size * 2 | |
| next_channels = size_to_channel[next_size] | |
| self.upsample_blocks.append(conv_block_factory.create_conv3_block( | |
| current_channels + next_channels, | |
| next_channels)) | |
| current_size = next_size | |
| current_channels = next_channels | |
| self.output_image_sizes.append(current_size) | |
| self.output_num_channels.append(current_channels) | |
| if args.upsample_mode == 'nearest': | |
| align_corners = None | |
| else: | |
| align_corners = False | |
| self.double_resolution = Upsample(scale_factor=2, mode=args.upsample_mode, align_corners=align_corners) | |
| def forward(self, feature: Tensor) -> List[Tensor]: | |
| downsampled_features = [] | |
| for block in self.downsample_blocks: | |
| feature = block(feature) | |
| downsampled_features.append(feature) | |
| for block in self.bottleneck_blocks: | |
| feature = block(feature) | |
| outputs = [feature] | |
| for i in range(0, len(self.upsample_blocks)): | |
| feature = self.double_resolution(feature) | |
| feature = torch.cat([feature, downsampled_features[-i - 2]], dim=1) | |
| feature = self.upsample_blocks[i](feature) | |
| outputs.append(feature) | |
| return outputs | |
| if __name__ == "__main__": | |
| device = torch.device('cuda') | |
| image_size = 512 | |
| image_channels = 4 | |
| num_pose_params = 6 | |
| args = ResizeConvUNetArgs( | |
| image_size=512, | |
| input_channels=10, | |
| start_channels=32, | |
| bottleneck_image_size=32, | |
| num_bottleneck_blocks=6, | |
| max_channels=512, | |
| upsample_mode='nearest', | |
| use_separable_convolution=False, | |
| block_args=BlockArgs( | |
| initialization_method='he', | |
| use_spectral_norm=False, | |
| normalization_layer_factory=InstanceNorm2dFactory(), | |
| nonlinearity_factory=ReLUFactory(inplace=False))) | |
| module = ResizeConvUNet(args).to(device) | |
| image_count = 8 | |
| input = torch.zeros(image_count, 10, 512, 512, device=device) | |
| outputs = module.forward(input) | |
| for output in outputs: | |
| print(output.shape) | |
| if True: | |
| repeat = 100 | |
| acc = 0.0 | |
| for i in range(repeat + 2): | |
| start = torch.cuda.Event(enable_timing=True) | |
| end = torch.cuda.Event(enable_timing=True) | |
| start.record() | |
| module.forward(input) | |
| end.record() | |
| torch.cuda.synchronize() | |
| if i >= 2: | |
| elapsed_time = start.elapsed_time(end) | |
| print("%d:" % i, elapsed_time) | |
| acc = acc + elapsed_time | |
| print("average:", acc / repeat) |