126 lines
4.0 KiB
Python
126 lines
4.0 KiB
Python
from typing import Optional
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from .channel_attn import CAResBlock
|
|
|
|
def interpolate_groups(g: torch.Tensor, ratio: float, mode: str,
|
|
align_corners: bool) -> torch.Tensor:
|
|
batch_size, num_objects = g.shape[:2]
|
|
g = F.interpolate(g.flatten(start_dim=0, end_dim=1),
|
|
scale_factor=ratio,
|
|
mode=mode,
|
|
align_corners=align_corners)
|
|
g = g.view(batch_size, num_objects, *g.shape[1:])
|
|
return g
|
|
|
|
|
|
def upsample_groups(g: torch.Tensor,
|
|
ratio: float = 2,
|
|
mode: str = 'bilinear',
|
|
align_corners: bool = False) -> torch.Tensor:
|
|
return interpolate_groups(g, ratio, mode, align_corners)
|
|
|
|
|
|
def downsample_groups(g: torch.Tensor,
|
|
ratio: float = 1 / 2,
|
|
mode: str = 'area',
|
|
align_corners: bool = None) -> torch.Tensor:
|
|
return interpolate_groups(g, ratio, mode, align_corners)
|
|
|
|
|
|
class GConv2d(nn.Conv2d):
|
|
def forward(self, g: torch.Tensor) -> torch.Tensor:
|
|
batch_size, num_objects = g.shape[:2]
|
|
g = super().forward(g.flatten(start_dim=0, end_dim=1))
|
|
return g.view(batch_size, num_objects, *g.shape[1:])
|
|
|
|
|
|
class GroupResBlock(nn.Module):
|
|
def __init__(self, in_dim: int, out_dim: int):
|
|
super().__init__()
|
|
|
|
if in_dim == out_dim:
|
|
self.downsample = nn.Identity()
|
|
else:
|
|
self.downsample = GConv2d(in_dim, out_dim, kernel_size=1)
|
|
|
|
self.conv1 = GConv2d(in_dim, out_dim, kernel_size=3, padding=1)
|
|
self.conv2 = GConv2d(out_dim, out_dim, kernel_size=3, padding=1)
|
|
|
|
def forward(self, g: torch.Tensor) -> torch.Tensor:
|
|
out_g = self.conv1(F.relu(g))
|
|
out_g = self.conv2(F.relu(out_g))
|
|
|
|
g = self.downsample(g)
|
|
|
|
return out_g + g
|
|
|
|
|
|
class MainToGroupDistributor(nn.Module):
|
|
def __init__(self,
|
|
x_transform: Optional[nn.Module] = None,
|
|
g_transform: Optional[nn.Module] = None,
|
|
method: str = 'cat',
|
|
reverse_order: bool = False):
|
|
super().__init__()
|
|
|
|
self.x_transform = x_transform
|
|
self.g_transform = g_transform
|
|
self.method = method
|
|
self.reverse_order = reverse_order
|
|
|
|
def forward(self, x: torch.Tensor, g: torch.Tensor, skip_expand: bool = False) -> torch.Tensor:
|
|
num_objects = g.shape[1]
|
|
|
|
if self.x_transform is not None:
|
|
x = self.x_transform(x)
|
|
|
|
if self.g_transform is not None:
|
|
g = self.g_transform(g)
|
|
|
|
if not skip_expand:
|
|
x = x.unsqueeze(1).expand(-1, num_objects, -1, -1, -1)
|
|
if self.method == 'cat':
|
|
if self.reverse_order:
|
|
g = torch.cat([g, x], 2)
|
|
else:
|
|
g = torch.cat([x, g], 2)
|
|
elif self.method == 'add':
|
|
g = x + g
|
|
elif self.method == 'mulcat':
|
|
g = torch.cat([x * g, g], dim=2)
|
|
elif self.method == 'muladd':
|
|
g = x * g + g
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
return g
|
|
|
|
|
|
class GroupFeatureFusionBlock(nn.Module):
|
|
def __init__(self, x_in_dim: int, g_in_dim: int, out_dim: int):
|
|
super().__init__()
|
|
|
|
x_transform = nn.Conv2d(x_in_dim, out_dim, kernel_size=1)
|
|
g_transform = GConv2d(g_in_dim, out_dim, kernel_size=1)
|
|
|
|
self.distributor = MainToGroupDistributor(x_transform=x_transform,
|
|
g_transform=g_transform,
|
|
method='add')
|
|
self.block1 = CAResBlock(out_dim, out_dim)
|
|
self.block2 = CAResBlock(out_dim, out_dim)
|
|
|
|
def forward(self, x: torch.Tensor, g: torch.Tensor) -> torch.Tensor:
|
|
batch_size, num_objects = g.shape[:2]
|
|
|
|
g = self.distributor(x, g)
|
|
|
|
g = g.flatten(start_dim=0, end_dim=1)
|
|
|
|
g = self.block1(g)
|
|
g = self.block2(g)
|
|
|
|
g = g.view(batch_size, num_objects, *g.shape[1:])
|
|
|
|
return g |